"""FastAPI app for the MedSegAgent chat-based segmentation demo.

Run from the repo root (so OAI_CONFIG_LIST/dataset/ paths in config.py and
plain package imports resolve correctly):

    cd /home/nlab/MedSegAgent
    source .venv/bin/activate
    uvicorn webapp.backend.main:app --host 0.0.0.0 --port 8000 --reload
"""

import logging
import shutil
import time
import uuid

from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from . import config, jobs
from .label_matcher import resolve_target
from .mask_extractor import extract_single_label
from .modality_detector import detect_modality
from .query_parser import parse_query
from .segmentation_runner import GPU_LOCK, run_totalsegmentator

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("medsegagent_webapp")

app = FastAPI(title="MedSegAgent Demo")

config.DATA_DIR.mkdir(parents=True, exist_ok=True)


class NoCacheStaticFiles(StaticFiles):
    """StaticFiles with caching disabled.

    Without this, browsers can skip revalidating index.html/app.js/style.css
    on a normal reload (heuristic caching kicks in when no Cache-Control
    header is set) and silently keep serving a stale build after an edit --
    a real footgun during active development of this demo. /static (the
    frontend source) always uses this; /data (per-job NIfTI results) does
    not need it since a given job_id's files never change after creation.
    """

    def file_response(self, *args, **kwargs):
        response = super().file_response(*args, **kwargs)
        response.headers["Cache-Control"] = "no-store"
        return response


app.mount("/data", StaticFiles(directory=str(config.DATA_DIR)), name="data")
app.mount("/static", NoCacheStaticFiles(directory=str(config.FRONTEND_DIR)), name="static")


@app.get("/")
def index():
    return FileResponse(str(config.FRONTEND_DIR / "index.html"), headers={"Cache-Control": "no-store"})


def _cleanup_old_job_dirs():
    """Delete job directories older than config.JOB_MAX_AGE_SECONDS.
    Called on every new /api/chat request; cheap for a handful of jobs."""
    now = time.time()
    if not config.DATA_DIR.exists():
        return
    for entry in config.DATA_DIR.iterdir():
        if entry.name == "samples" or not entry.is_dir():
            continue
        try:
            if now - entry.stat().st_mtime > config.JOB_MAX_AGE_SECONDS:
                shutil.rmtree(entry, ignore_errors=True)
        except FileNotFoundError:
            pass


def _run_segmentation_job(job_id, input_path, seg_path, mask_path, modality, label_id):
    """Background task: waits for the GPU lock, runs inference, extracts just
    the target label into its own file, and updates job status.

    Status transitions: queued -> running -> done | error. The "queued"
    state persists (accurately) while waiting on GPU_LOCK for a second
    concurrent request, not just while waiting for FastAPI to schedule the
    background task.

    seg_path ends up holding TotalSegmentator's full multi-label output
    (every structure it found); mask_path holds a derived single-label mask
    containing only `label_id` -- see mask_extractor.py for why the frontend
    is served mask_path, not seg_path, as the overlay.
    """
    try:
        jobs.update_job(job_id, status="queued", message="GPUの空き待ちです...")
        with GPU_LOCK:
            jobs.update_job(job_id, status="running", message="セグメンテーション実行中...")
            run_totalsegmentator(input_path, seg_path, modality)
        extract_single_label(seg_path, label_id, mask_path)
        jobs.update_job(job_id, status="done", message="完了しました。")
    except Exception as e:  # noqa: BLE001 - surface any failure to the UI
        logger.exception("segmentation job %s failed", job_id)
        jobs.update_job(job_id, status="error", message=f"推論に失敗しました: {e}")


@app.post("/api/chat")
async def chat(
    background_tasks: BackgroundTasks,
    query: str = Form(...),
    file: UploadFile = File(...),
):
    _cleanup_old_job_dirs()

    filename = file.filename or ""
    if not filename.lower().endswith(config.ALLOWED_UPLOAD_EXTENSIONS):
        raise HTTPException(
            status_code=400,
            detail="対応していないファイル形式です（.nii.gz または .nii のみアップロードできます）。",
        )

    job_id = uuid.uuid4().hex
    job_dir = config.DATA_DIR / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    input_path = job_dir / "input.nii.gz"

    size = 0
    try:
        with open(input_path, "wb") as f:
            while True:
                chunk = await file.read(1024 * 1024)
                if not chunk:
                    break
                size += len(chunk)
                if size > config.MAX_UPLOAD_BYTES:
                    raise HTTPException(
                        status_code=413,
                        detail="ファイルサイズが大きすぎます（上限1GB）。",
                    )
                f.write(chunk)
    except HTTPException:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise

    try:
        parsed = parse_query(query)
    except Exception as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(status_code=400, detail=f"クエリを解析できませんでした: {e}")

    target = parsed["target"]

    # Modality routing is decided from the uploaded volume's own voxel
    # intensities, NOT from the LLM's parsed query text -- a query like
    # "segment the liver" very often won't mention the modality at all even
    # though the file itself is unambiguously a CT or MRI. See
    # modality_detector.py for why. parsed["modality"] (e.g. "not specified",
    # or whatever the user's wording implied) is kept only for display.
    try:
        modality = detect_modality(input_path)
    except Exception as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(status_code=400, detail=f"画像のモダリティを判定できませんでした: {e}")

    try:
        resolved_label = resolve_target(target, modality, model=config.MODEL_NAME)
    except Exception as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(status_code=400, detail=str(e))

    seg_path = job_dir / "segmentation.nii.gz"
    mask_path = job_dir / "target_mask.nii.gz"
    jobs.create_job(
        job_id,
        status="queued",
        message="GPUの空き待ちです...",
        modality=modality,
        modality_from_query=parsed.get("modality"),
        target=target,
        resolved_label=resolved_label,
    )
    background_tasks.add_task(
        _run_segmentation_job,
        job_id,
        input_path,
        seg_path,
        mask_path,
        modality,
        resolved_label["id"],
    )

    return {
        "job_id": job_id,
        "modality": modality,
        "modality_from_query": parsed.get("modality"),
        "target": target,
        "resolved_label": resolved_label,
        "status": "queued",
    }


@app.get("/api/jobs/{job_id}")
def get_job_status(job_id: str):
    job = jobs.get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail="ジョブが見つかりません。")

    result = None
    if job["status"] == "done":
        result = {
            "input_url": f"/data/{job_id}/input.nii.gz",
            # target_mask.nii.gz (not the raw multi-label segmentation.nii.gz)
            # -- contains only 0 and label_id, see mask_extractor.py.
            "seg_url": f"/data/{job_id}/target_mask.nii.gz",
            "label_id": job["resolved_label"]["id"],
            "label_name": job["resolved_label"]["name"],
        }

    return {
        "status": job["status"],
        "message": job.get("message", ""),
        "modality": job.get("modality"),
        "target": job.get("target"),
        "resolved_label": job.get("resolved_label"),
        "result": result,
    }