stage_read/api.py

"""FastAPI router for Stage Read jobs.

POST /stage-read                 multipart: file, options (json string)   -> {job_id}
GET  /stage-read/{job_id}        -> {status, stage, percent, message, result?, error?}
GET  /stage-read/{job_id}/sheet/{speaker}  -> JPEG contact sheet for naming a speaker label

Jobs run in a background thread, one at a time (the GPU is the bottleneck). State is in memory;
the working directory keeps every intermediate file so a restarted service can resume a job by
re-posting the same file with the same job_id.

RETENTION. A job's folder holds the uploaded recording, face contact sheets and the result, so it
is deleted STAGE_READ_TTL_HOURS (default 24) after it was created: long enough to name speakers
and save the report in the app, which keeps its own copy of the result. Swept on every submit.
Every request also needs the app's signed pass (service_auth.py).
"""
import os, re, json, uuid, time, threading, tempfile, shutil, traceback
from fastapi import APIRouter, UploadFile, File, Form, HTTPException
from fastapi.responses import FileResponse
from . import pipeline

router = APIRouter()
JOBS: dict = {}
LOCK = threading.Lock()
WORK_ROOT = os.environ.get("STAGE_READ_WORK", os.path.join(tempfile.gettempdir(), "stage_read"))
TTL_S = float(os.environ.get("STAGE_READ_TTL_HOURS", "24")) * 3600

# Both end up in file paths, so nothing but the shapes we mint gets through
JOB_ID = re.compile(r"^[0-9a-f]{32}$")
SPEAKER = re.compile(r"^[A-Za-z0-9_-]{1,40}$")


def _job_dir(job_id: str) -> str:
    if not JOB_ID.match(job_id or ""):
        raise HTTPException(400, "bad job id")
    return os.path.join(WORK_ROOT, job_id)


def _check_speaker(speaker: str) -> None:
    if not SPEAKER.match(speaker or ""):
        raise HTTPException(400, "bad speaker")


def _sweep() -> None:
    """Delete job folders past their retention, and forget their in-memory state."""
    if not os.path.isdir(WORK_ROOT):
        return
    cutoff = time.time() - TTL_S
    for name in os.listdir(WORK_ROOT):
        path = os.path.join(WORK_ROOT, name)
        try:
            if os.path.isdir(path) and os.path.getmtime(path) < cutoff and not JOBS.get(name, {}).get("running"):
                shutil.rmtree(path, ignore_errors=True)
                JOBS.pop(name, None)
        except OSError:
            pass


def _run(job_id: str, video: str, options: dict):
    job = JOBS[job_id]
    def progress(stage, pct, msg):
        job.update(status="running", stage=stage, percent=pct, message=msg)
    try:
        workdir = os.path.join(WORK_ROOT, job_id)
        result = pipeline.run(video, workdir, options, progress)
        job["sheets"] = pipeline.contact_sheet(video, workdir, result, options.get("panel"))
        job.update(status="done", result=result, percent=100, stage="done", message="Complete")
    except Exception as e:
        job.update(status="error", error=f"{e}", trace=traceback.format_exc())
    finally:
        with LOCK:
            job["running"] = False


@router.post("/stage-read")
def submit(file: UploadFile = File(...), options: str = Form("{}"), job_id: str = Form(None)):
    _sweep()
    opts = json.loads(options or "{}")
    job_id = job_id or uuid.uuid4().hex
    workdir = _job_dir(job_id); os.makedirs(workdir, exist_ok=True)
    video = os.path.join(workdir, "input" + os.path.splitext(file.filename or ".mp4")[1])
    with open(video, "wb") as f:
        shutil.copyfileobj(file.file, f)
    # Beside the intermediate files, so a resume after the service restarts measures the same job
    # rather than silently reverting to the default panel and models.
    with open(os.path.join(workdir, "options.json"), "w", encoding="utf-8") as f:
        json.dump(opts, f)
    with LOCK:
        if any(j.get("running") for j in JOBS.values()):
            JOBS[job_id] = {"status": "queued", "stage": "queued", "percent": 0, "message": "Another job is running", "running": False, "video": video, "options": opts}
            return {"job_id": job_id, "status": "queued"}
        JOBS[job_id] = {"status": "running", "stage": "audio", "percent": 0, "message": "Starting", "running": True, "video": video, "options": opts}
    threading.Thread(target=_run, args=(job_id, video, opts), daemon=True).start()
    return {"job_id": job_id, "status": "running"}


@router.get("/stage-read/{job_id}")
def status(job_id: str):
    workdir = _job_dir(job_id)
    job = JOBS.get(job_id)
    if not job:
        # a finished job from a previous process: serve its file if it exists
        p = os.path.join(workdir, "stage_read.json")
        if os.path.exists(p):
            return {"status": "done", "stage": "done", "percent": 100, "message": "Complete", "result": json.load(open(p, encoding="utf-8"))}
        raise HTTPException(404, "unknown job")
    with LOCK:
        if job["status"] == "queued" and not any(j.get("running") for j in JOBS.values()):
            job.update(status="running", running=True)
            threading.Thread(target=_run, args=(job_id, job["video"], job["options"]), daemon=True).start()
    out = {k: job.get(k) for k in ("status", "stage", "percent", "message", "error")}
    if job["status"] == "done":
        out["result"] = job["result"]; out["sheets"] = sorted(job.get("sheets", {}).keys())
    return out


@router.post("/stage-read/{job_id}/resume")
def resume(job_id: str):
    """Restart a job whose process died mid-run. Cached steps in the working directory are reused."""
    workdir = _job_dir(job_id)
    inputs = [f for f in os.listdir(workdir)] if os.path.isdir(workdir) else []
    video = next((os.path.join(workdir, f) for f in inputs if f.startswith("input")), None)
    if not video:
        raise HTTPException(404, "no input for that job")
    opts_path = os.path.join(workdir, "options.json")
    opts = json.load(open(opts_path, encoding="utf-8")) if os.path.exists(opts_path) else {}
    with LOCK:
        if any(j.get("running") for j in JOBS.values()):
            raise HTTPException(409, "another job is running")
        JOBS[job_id] = {"status": "running", "stage": "resume", "percent": 0, "message": "Resuming", "running": True, "video": video, "options": opts}
    threading.Thread(target=_run, args=(job_id, video, opts), daemon=True).start()
    return {"job_id": job_id, "status": "running"}


@router.get("/stage-read/{job_id}/sheet/{speaker}")
def sheet(job_id: str, speaker: str):
    _check_speaker(speaker)
    p = os.path.join(_job_dir(job_id), "sheets", f"{speaker}.jpg")
    if not os.path.exists(p):
        raise HTTPException(404, "no sheet")
    return FileResponse(p, media_type="image/jpeg")


@router.get("/stage-read/{job_id}/plate/{speaker}")
def plate(job_id: str, speaker: str, when: str = "intro"):
    """High-resolution crop of the slide as the speaker was introduced, for reading the name card.

    Built on demand rather than during the run: it is only wanted while someone is naming speakers,
    it costs a few seeks over a video already on disk, and generating it here means finished jobs
    get it without being re-run. The result is cached beside the contact sheets.
    """
    _check_speaker(speaker)
    workdir = _job_dir(job_id)
    result_path = os.path.join(workdir, "stage_read.json")
    if not os.path.exists(result_path):
        raise HTTPException(404, "no result for that job")
    video = next((os.path.join(workdir, f) for f in os.listdir(workdir) if f.startswith("input")), None)
    if not video:
        raise HTTPException(404, "no input for that job")
    if when not in ("intro", "turn"):
        raise HTTPException(400, "when must be intro or turn")
    result = json.load(open(result_path, encoding="utf-8"))
    try:
        p = pipeline.nameplate(video, workdir, result, speaker, when)
    except Exception as e:
        raise HTTPException(500, f"nameplate failed: {e}")
    if not p:
        raise HTTPException(404, "no frame for that speaker")
    return FileResponse(p, media_type="image/jpeg")