"""FastAPI application: uploads, validation, jobs, results, examples and the web UI."""
from __future__ import annotations

import asyncio
import json
import mimetypes
import re
import shutil
import threading
import time
import uuid
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse, Response, StreamingResponse
from pydantic import BaseModel, Field

from ppi_controls import PRESETS, Inputs, validate
from ppi_controls.constants import ENGINE_VERSION, FRAMEWORK_VERSION
from ppi_controls.io import InputError, read_pairs_with_issues, read_proteins_with_issues, sniff

from . import db
from .formatspec import FORMAT_SPEC, template
from .settings import settings
from .worker import JobManager

# Slim images ship without /etc/mime.types; the help page serves these from static/.
mimetypes.add_type("image/webp", ".webp")
mimetypes.add_type("application/vnd.openxmlformats-officedocument.wordprocessingml.document", ".docx")

ID_RE = re.compile(r"^[0-9a-f]{32}$")
FILE_KINDS = {
    "proteins": {".tsv", ".csv", ".txt", ".tab", ".fa", ".fasta", ".faa"},
    "pairs": {".tsv", ".csv", ".txt", ".tab"},
    "species": {".tsv", ".csv", ".txt", ".tab"},
    "clusters": {".tsv", ".txt", ".tab"},
    "embeddings": {".csv", ".tsv", ".npz", ".h5", ".hdf5"},
}
TEXT_KINDS = {"proteins", "pairs", "species"}

STATE: dict = {}


def _purge_loop():
    """Retention: remove datasets and finished jobs older than the configured age."""
    while True:
        try:
            cutoff = time.time() - settings.retention_days * 86400
            for job_id in db.old_jobs(settings.db_path, cutoff):
                shutil.rmtree(settings.storage / "jobs" / job_id, ignore_errors=True)
                db.delete(settings.db_path, job_id)
            for d in (settings.storage / "datasets").iterdir():
                if d.is_dir() and d.stat().st_mtime < cutoff:
                    shutil.rmtree(d, ignore_errors=True)
        except Exception:  # noqa: BLE001 - never let the janitor kill the server
            pass
        time.sleep(3600)


@asynccontextmanager
async def lifespan(app: FastAPI):
    settings.ensure()
    db.init(settings.db_path)
    STATE["jobs"] = JobManager(settings)
    threading.Thread(target=_purge_loop, daemon=True).start()
    yield
    STATE["jobs"].shutdown()


app = FastAPI(title="PPI Benchmark Studio", version=ENGINE_VERSION, lifespan=lifespan)
if settings.cors_origins:
    app.add_middleware(CORSMiddleware, allow_origins=list(settings.cors_origins), allow_methods=["*"], allow_headers=["*"])


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _check_id(value: str) -> str:
    if not ID_RE.match(value or ""):
        raise HTTPException(404, "Not found")
    return value


def _dataset_dir(dataset_id: str) -> Path:
    d = settings.storage / "datasets" / _check_id(dataset_id)
    if not (d / "meta.json").exists():
        raise HTTPException(404, "Dataset not found or expired")
    return d


def _job_dir(job_id: str) -> Path:
    return settings.storage / "jobs" / _check_id(job_id)


def _meta(d: Path) -> dict:
    return json.loads((d / "meta.json").read_text())


def _issue_error(status: int, message: str, issues=None):
    return JSONResponse(status_code=status, content={"detail": message, "issues": issues or []})


async def _save_upload(upload: UploadFile, dest: Path) -> int:
    limit = settings.max_upload_mb * 1024 * 1024
    size = 0
    with dest.open("wb") as fh:
        while chunk := await upload.read(1 << 20):
            size += len(chunk)
            if size > limit:
                fh.close()
                dest.unlink(missing_ok=True)
                raise HTTPException(413, f"{upload.filename} exceeds the {settings.max_upload_mb} MB upload limit")
            fh.write(chunk)
    return size


def _job_view(job: dict) -> dict:
    view = {k: job[k] for k in ("id", "dataset_id", "status", "stage", "stage_label", "progress", "message",
                                "config", "error", "created", "started", "finished")}
    if job["status"] == "queued":
        view["queue_position"] = db.queue_position(settings.db_path, job["id"])
    return view


# ---------------------------------------------------------------------------
# meta
# ---------------------------------------------------------------------------

@app.get("/api/health")
def health():
    return {"status": "ok", "engine_version": ENGINE_VERSION, "framework_version": FRAMEWORK_VERSION}


@app.get("/api/config")
def public_config():
    return {"presets": {k: v for k, v in PRESETS.items() if k in settings.allowed_presets},
            "limits": {"max_upload_mb": settings.max_upload_mb, "max_pairs": settings.max_pairs,
                       "retention_days": settings.retention_days, "max_concurrent_jobs": settings.max_concurrent_jobs},
            "engine_version": ENGINE_VERSION, "framework_version": FRAMEWORK_VERSION}


@app.get("/api/format")
def format_spec():
    return {**FORMAT_SPEC, "limits": {"max_upload_mb": settings.max_upload_mb, "max_pairs": settings.max_pairs}}


@app.get("/api/templates/{name}")
def get_template(name: str):
    t = template(name)
    if t is None:
        raise HTTPException(404, "Unknown template")
    content, media = t
    return Response(content, media_type=media, headers={"Content-Disposition": f'attachment; filename="{name}"'})


class CheckRequest(BaseModel):
    kind: str = Field(pattern="^(proteins|pairs|species)$")
    text: str = Field(max_length=200_000)


@app.post("/api/check")
def check_snippet(req: CheckRequest):
    """Format checker for pasted text (format-guide page)."""
    sn = sniff(req.text, req.kind)
    issues, rows = [], 0
    try:
        if req.kind == "proteins":
            df, issues = read_proteins_with_issues(req.text)
        elif req.kind == "pairs":
            df, issues = read_pairs_with_issues(req.text)
        else:
            df = None
        rows = 0 if df is None else len(df)
    except InputError as e:
        issues = e.issues
    return {"sniff": sn.to_dict(), "rows_parsed": rows, "issues": [i.to_dict() for i in issues],
            "ok": not any(i.severity == "error" for i in issues)}


# ---------------------------------------------------------------------------
# datasets
# ---------------------------------------------------------------------------

@app.post("/api/datasets")
async def create_dataset(proteins: UploadFile = File(...), pairs: UploadFile = File(...),
                         species: UploadFile | None = File(None), clusters: UploadFile | None = File(None),
                         embeddings: UploadFile | None = File(None)):
    dataset_id = uuid.uuid4().hex
    d = settings.storage / "datasets" / dataset_id
    d.mkdir(parents=True)
    files, previews = {}, {}
    try:
        for kind, up in (("proteins", proteins), ("pairs", pairs), ("species", species), ("clusters", clusters),
                         ("embeddings", embeddings)):
            if up is None or not up.filename:
                continue
            ext = Path(up.filename).suffix.lower()
            if ext not in FILE_KINDS[kind]:
                raise HTTPException(415, f"{kind}: unsupported file type {ext or '(none)'}; use {sorted(FILE_KINDS[kind])}")
            stored = f"{kind}{ext}"
            size = await _save_upload(up, d / stored)
            files[kind] = {"name": Path(up.filename).name[:200], "stored": stored, "bytes": size}
            if kind in TEXT_KINDS:
                previews[kind] = sniff(d / stored, kind).to_dict()
        if previews["pairs"]["n_rows"] > settings.max_pairs:
            raise HTTPException(413, f"{previews['pairs']['n_rows']:,} pairs exceed the {settings.max_pairs:,}-pair limit "
                                     "of this server; run the local Docker version for larger datasets")
    except HTTPException:
        shutil.rmtree(d, ignore_errors=True)
        raise
    meta = {"id": dataset_id, "created": time.time(), "files": files, "previews": previews, "mappings": {}}
    (d / "meta.json").write_text(json.dumps(meta))
    return meta


@app.get("/api/datasets/{dataset_id}")
def get_dataset(dataset_id: str):
    return _meta(_dataset_dir(dataset_id))


class ValidateRequest(BaseModel):
    mappings: dict = Field(default_factory=dict)


@app.post("/api/datasets/{dataset_id}/validate")
def validate_dataset(dataset_id: str, req: ValidateRequest | None = None):
    d = _dataset_dir(dataset_id)
    meta = _meta(d)
    if req and req.mappings:
        meta["mappings"] = {k: v for k, v in req.mappings.items() if k in TEXT_KINDS and isinstance(v, dict)}
        (d / "meta.json").write_text(json.dumps(meta))
    f = {k: d / v["stored"] for k, v in meta["files"].items()}
    try:
        report = validate(Inputs(f["proteins"], f["pairs"], f.get("species"), mappings=meta["mappings"]))
    except InputError as e:
        report = {"ok": False, "issues": [i.to_dict() for i in e.issues], "audit": None}
    meta["validated"] = report["ok"]
    (d / "meta.json").write_text(json.dumps(meta))
    return report


# ---------------------------------------------------------------------------
# jobs
# ---------------------------------------------------------------------------

class JobRequest(BaseModel):
    dataset_id: str
    preset: str = "quick"
    seed: int = Field(42, ge=0, le=2**31 - 1)
    species_mode: str = Field("auto", pattern="^(auto|off)$")
    embeddings_name: str | None = Field(None, max_length=60)


@app.post("/api/jobs", status_code=202)
def create_job(req: JobRequest):
    d = _dataset_dir(req.dataset_id)
    if req.preset not in settings.allowed_presets or req.preset not in PRESETS:
        raise HTTPException(400, f"Preset must be one of {list(settings.allowed_presets)}")
    if db.count(settings.db_path, ("queued",)) >= settings.max_queued_jobs:
        raise HTTPException(429, "The server queue is full; please try again in a few minutes")
    job_id = uuid.uuid4().hex
    job_dir = _job_dir(job_id)
    job_dir.mkdir(parents=True)
    config = req.model_dump(exclude={"dataset_id"})
    db.create_job(settings.db_path, job_id, req.dataset_id, config)
    STATE["jobs"].submit(job_id, d, job_dir, config)
    return _job_view(db.get(settings.db_path, job_id))


@app.get("/api/jobs/{job_id}")
def get_job(job_id: str):
    job = db.get(settings.db_path, _check_id(job_id))
    if job is None:
        raise HTTPException(404, "Job not found or expired")
    return _job_view(job)


@app.delete("/api/jobs/{job_id}")
def cancel_job(job_id: str):
    job = db.get(settings.db_path, _check_id(job_id))
    if job is None:
        raise HTTPException(404, "Job not found")
    if job["status"] not in db.TERMINAL:
        db.update(settings.db_path, job_id, cancel=1)
        if job["status"] == "queued":
            db.update(settings.db_path, job_id, status="cancelled", message="Cancelled", finished=time.time())
    return _job_view(db.get(settings.db_path, job_id))


@app.get("/api/jobs/{job_id}/events")
async def job_events(job_id: str, request: Request):
    _check_id(job_id)
    if db.get(settings.db_path, job_id) is None:
        raise HTTPException(404, "Job not found")

    async def stream():
        last = None
        while True:
            if await request.is_disconnected():
                return
            job = db.get(settings.db_path, job_id)
            if job is None:
                return
            view = _job_view(job)
            payload = json.dumps(view)
            if payload != last:
                last = payload
                yield f"data: {payload}\n\n"
            if job["status"] in db.TERMINAL:
                return
            await asyncio.sleep(0.5)
    return StreamingResponse(stream(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


def _finished_output(job_id: str) -> Path:
    job = db.get(settings.db_path, _check_id(job_id))
    if job is None:
        raise HTTPException(404, "Job not found or expired")
    if job["status"] != "done":
        raise HTTPException(409, f"Job is {job['status']}")
    return _job_dir(job_id) / "output"


@app.get("/api/jobs/{job_id}/results")
def job_results(job_id: str):
    return FileResponse(_finished_output(job_id) / "results.json", media_type="application/json")


@app.get("/api/jobs/{job_id}/bundle")
def job_bundle(job_id: str):
    _finished_output(job_id)
    return FileResponse(_job_dir(job_id) / "bundle.zip", media_type="application/zip",
                        filename=f"ppi_evaluation_{job_id[:8]}.zip")


def _safe_file(root: Path, rel: str) -> Path:
    path = (root / rel).resolve()
    if not path.is_file() or root.resolve() not in path.parents:
        raise HTTPException(404, "File not found")
    return path


@app.get("/api/jobs/{job_id}/files/{rel:path}")
def job_file(job_id: str, rel: str):
    out = _finished_output(job_id)
    return FileResponse(_safe_file(out, rel), filename=Path(rel).name)


# ---------------------------------------------------------------------------
# examples
# ---------------------------------------------------------------------------

def _examples() -> list[dict]:
    rows = []
    for d in sorted(settings.examples.glob("*/meta.json")):
        meta = json.loads(d.read_text())
        meta["ready"] = (d.parent / "output" / "results.json").exists()
        rows.append(meta)
    order = {"multi_species": 0, "guo_yeast": 1, "synthetic": 2}
    return sorted(rows, key=lambda m: order.get(m["name"], 9))


def _example_dir(name: str) -> Path:
    if name not in {m["name"] for m in _examples()}:
        raise HTTPException(404, "Unknown example")
    return settings.examples / name


@app.get("/api/examples")
def list_examples():
    return _examples()


@app.get("/api/examples/{name}/results")
def example_results(name: str):
    path = _example_dir(name) / "output" / "results.json"
    if not path.exists():
        raise HTTPException(409, "Example results are not built yet")
    return FileResponse(path, media_type="application/json")


@app.get("/api/examples/{name}/files/{rel:path}")
def example_file(name: str, rel: str):
    return FileResponse(_safe_file(_example_dir(name), rel), filename=Path(rel).name)


@app.post("/api/examples/{name}/dataset")
def example_as_dataset(name: str):
    """Copy an example's input files into a new dataset, to try the upload → run flow."""
    src = _example_dir(name) / "input"
    dataset_id = uuid.uuid4().hex
    d = settings.storage / "datasets" / dataset_id
    d.mkdir(parents=True)
    files, previews = {}, {}
    for kind in ("proteins", "pairs", "species"):
        f = src / f"{kind}.tsv"
        if f.exists():
            shutil.copy(f, d / f.name)
            files[kind] = {"name": f.name, "stored": f.name, "bytes": f.stat().st_size}
            previews[kind] = sniff(d / f.name, kind).to_dict()
    meta = {"id": dataset_id, "created": time.time(), "files": files, "previews": previews, "mappings": {}, "example": name}
    (d / "meta.json").write_text(json.dumps(meta))
    return meta


# ---------------------------------------------------------------------------
# web UI (built frontend), with single-page-app fallback
# ---------------------------------------------------------------------------

@app.get("/{rel:path}", include_in_schema=False)
def spa(rel: str):
    if rel.startswith("api/"):
        raise HTTPException(404, "Not found")
    static = settings.static
    if rel:
        candidate = (static / rel).resolve()
        if candidate.is_file() and static.resolve() in candidate.parents:
            return FileResponse(candidate)
    index = static / "index.html"
    if index.exists():
        return FileResponse(index)
    return PlainTextResponse("PPI Benchmark Studio API is running. Build the frontend (frontend/) to serve the web UI.")


def serve():
    import argparse

    import uvicorn
    ap = argparse.ArgumentParser(prog="ppi-studio-api")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8000)
    a = ap.parse_args()
    uvicorn.run("studio_api.main:app", host=a.host, port=a.port, proxy_headers=True)