app.py

import base64
import io
import tempfile
import uuid
from pathlib import Path

from gradio import Server
from fastapi import Request
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse

from commentary import STYLES, generate_commentary          # eager-loads Qwen on import
from events import events_to_frames, pick_key_events
from faces import annotate_event_frame, cluster_faces, frame_at, scan_video
from ov_models import download_models
from video import get_duration, normalize_video

MAX_SECONDS = 60.0          # keep the demo snappy; matches the UI copy
HERE = Path(__file__).parent
SESSIONS = {}

TIERS = ["HEADLINER", "RISING", "DIVA", "WILDCARD", "BREAKOUT", "CAMEO"]
VIBE_TO_STYLE = {
    "football": "Football (Premier League hype)",
    "diva":     "Diva Hour",
    "wildlife": "Nature documentary",
    "boxing":   "Boxing announcer",
}
SUPPORTED_VIDEO = {".mp4", ".mov", ".webm", ".mkv", ".avi", ".m4v"}

try:
    download_models()
except Exception as e:
    print(f"[warn] OV model download failed at boot: {e}")

def _crop_b64(pil, size=160):
    img = pil.copy(); img.thumbnail((size, size))
    buf = io.BytesIO(); img.save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()

def _build_ui() -> str:
    """Inline CSS + JS into Marquee.html so the page is fully self-contained."""
    doc = (HERE / "Marquee.html").read_text()
    css = (HERE / "marquee.css").read_text()
    js  = (HERE / "marquee.js").read_text()

    doc = doc.replace('<link rel="stylesheet" href="marquee.css">',
                      f"<style>{css}</style>")
    doc = doc.replace('<script src="marquee.js"></script>',
                      f"<script type='module'>{js}</script>")
    return doc

app = Server()

_UI_HTML = _build_ui()


@app.get("/", response_class=HTMLResponse)
async def root():
    return HTMLResponse(_UI_HTML)


@app.post("/api/scan")
async def api_scan(request: Request):
    """CPU-only (OpenVINO + ffmpeg). Plain route + multipart upload via fetch."""
    try:
        form = await request.form()
        upload = form.get("file")
        if upload is None:
            return JSONResponse({"ok": False, "error": "no file"}, status_code=400)

        suffix = Path(upload.filename).suffix.lower()
        if suffix not in SUPPORTED_VIDEO:
            return JSONResponse({"ok": False,
                "error": f"Unsupported format '{suffix}'. Supported: {', '.join(sorted(SUPPORTED_VIDEO))}"})

        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
        tmp.write(await upload.read()); tmp.flush(); tmp.close()

        dur = get_duration(tmp.name)
        if dur > MAX_SECONDS + 1:
            return JSONResponse({"ok": False, "error": "Clip longer than 60 seconds"})

        norm_path, w, h = normalize_video(tmp.name)
        embeds, crops, motion, meta = scan_video(norm_path, MAX_SECONDS)
        identities = cluster_faces(embeds, crops)

        sid = uuid.uuid4().hex
        SESSIONS[sid] = {"norm": norm_path, "identities": identities,
                         "motion": motion, "meta": meta}

        stars = [{"idx": i, "tier": TIERS[i % len(TIERS)],
                  "crop": _crop_b64(idn["crop"])}
                 for i, idn in enumerate(identities)]

        return JSONResponse({"ok": True, "session_id": sid,
            "video_url": f"/video/{sid}",
            "duration": get_duration(norm_path),
            "w": w, "h": h, "portrait": h > w, "stars": stars})

    except Exception as e:
        import traceback; traceback.print_exc()
        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)


@app.get("/video/{sid}")
async def serve_video(sid: str):
    """Stream the normalized video file."""
    sess = SESSIONS.get(sid)
    if not sess:
        return JSONResponse({"error": "session not found"}, status_code=404)
    return FileResponse(sess["norm"], media_type="video/mp4",
                        headers={"Accept-Ranges": "bytes"})


@app.api(name="generate")
def generate(session_id: str, names: list, vibe: str) -> dict:
    """GPU path. Exposed via @app.api() so ZeroGPU detects the @spaces.GPU
    function and gradio queues the calls. Called from the frontend through the
    Gradio JS client.
    """
    sess = SESSIONS.get(session_id)
    if not sess:
        return {"ok": False, "error": "Session expired -- re-cast"}

    name_map = {i: str(n).strip() for i, n in enumerate(names) if str(n).strip()}
    style    = VIBE_TO_STYLE.get(vibe, next(iter(STYLES)))

    fps         = sess["meta"]["fps"]
    event_times = pick_key_events(sess["motion"], fps)
    event_idxs  = events_to_frames(event_times, fps)

    frames = []
    for ts, fidx in zip(event_times, event_idxs):
        bgr = frame_at(sess["norm"], fidx)
        if bgr is None:
            continue
        frames.append((ts, annotate_event_frame(bgr, sess["identities"], name_map)))

    script = generate_commentary(frames, list(name_map.values()), style)
    return {"ok": True,
            "script": [{"t": round(float(l["time"]), 2), "text": l["text"]}
                       for l in script]}


# Custom "/" already takes priority over Gradio's default page; no gradio UI
# components, so SSR is pointless -> off. On Spaces, launch() binds 0.0.0.0:7860.
app.launch(ssr_mode=False, show_error=True)