app.py

import base64
import io
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from pathlib import Path

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Install chatterbox-tts with --no-deps BEFORE any other import.
# ---------------------------------------------------------------------------
try:
    import chatterbox  # noqa: F401
except ImportError:
    log.info("[boot] Installing chatterbox-tts --no-deps …")
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-deps", "--quiet", "chatterbox-tts"],
        check=True,
    )
    log.info("[boot] chatterbox-tts installed.")

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

from commentary import STYLES, VIBE_TO_STYLE, generate_commentary
from events import events_to_frames, pick_key_events
from faces import annotate_event_frame, frame_at, scan_video, write_names_to_jsonl
from ov_models import download_models
from tts import generate_script_audio
from video import get_duration, normalize_video

MAX_SECONDS = 60.0
HERE        = Path(__file__).parent
SESSIONS: dict[str, dict] = {}

TIERS = ["HEADLINER", "RISING", "DIVA", "WILDCARD", "BREAKOUT", "CAMEO"]
SUPPORTED_VIDEO = {".mp4", ".mov", ".webm", ".mkv", ".avi", ".m4v"}

log.info("[boot] Downloading OpenVINO models if needed…")
try:
    download_models()
    log.info("[boot] OpenVINO models ready.")
except Exception as e:
    log.warning(f"[boot] OV model download failed: {e}")


def _crop_b64(pil, size=160):
    if pil is None:
        return None
    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:
    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()
log.info("[boot] UI built and ready.")


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


@app.post("/api/scan")
async def api_scan(request: Request):
    t0 = time.time()
    log.info("[scan] Request received")
    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()
        log.info(f"[scan] File: {upload.filename} ({suffix})")
        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)
        data = await upload.read()
        tmp.write(data); tmp.flush(); tmp.close()
        log.info(f"[scan] Upload saved: {len(data)//1024}KB → {tmp.name}")

        dur = get_duration(tmp.name)
        log.info(f"[scan] Duration: {dur:.1f}s")
        if dur > MAX_SECONDS + 1:
            return JSONResponse({"ok": False, "error": "Clip longer than 60 seconds"})

        log.info("[scan] Normalising video (ffmpeg)…")
        t1 = time.time()
        norm_path, w, h = normalize_video(tmp.name)
        log.info(f"[scan] Normalised in {time.time()-t1:.1f}s → {w}x{h}")

        log.info("[scan] Running person tracking scan…")
        tracks, motion, meta, jsonl_path = scan_video(norm_path, MAX_SECONDS)
        log.info(f"[scan] Tracking done — {len(tracks)} person(s) found")

        sid = uuid.uuid4().hex
        SESSIONS[sid] = {
            "norm":       norm_path,
            "tracks":     tracks,
            "jsonl_path": jsonl_path,
            "motion":     motion,
            "meta":       meta,
        }
        log.info(f"[scan] Session created: {sid}")

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

        log.info(f"[scan] Completed in {time.time()-t0:.1f}s total")
        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()
        log.error(f"[scan] FAILED: {e}")
        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)


@app.get("/video/{sid}")
async def serve_video(sid: str):
    log.info(f"[video] Serving video for session {sid}")
    sess = SESSIONS.get(sid)
    if not sess:
        log.warning(f"[video] Session not found: {sid}")
        return JSONResponse({"error": "session not found"}, status_code=404)
    return FileResponse(sess["norm"], media_type="video/mp4",
                        headers={"Accept-Ranges": "bytes"})


@app.post("/api/tts")
async def api_tts(request: Request):
    t0 = time.time()
    log.info("[tts] Request received")
    try:
        body        = await request.json()
        sid         = body.get("session_id", "")
        vibe        = body.get("vibe", "football")
        script      = body.get("script", [])
        voice_model = body.get("voice_model", "chatterbox")

        log.info(f"[tts] sid={sid} vibe={vibe} voice={voice_model} lines={len(script)}")

        if not script:
            return JSONResponse({"ok": False, "error": "empty script"}, status_code=400)
        if sid and sid not in SESSIONS:
            return JSONResponse({"ok": False, "error": "session not found"}, status_code=404)

        log.info(f"[tts] Generating {len(script)} lines with {voice_model}…")
        lines = generate_script_audio(script, vibe, voice_model=voice_model)
        audio_count = sum(1 for l in lines if l.get("audio_b64"))
        log.info(f"[tts] Done in {time.time()-t0:.1f}s — "
                 f"{audio_count}/{len(lines)} lines have audio")
        return JSONResponse({"ok": True, "lines": lines})

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


@app.post("/api/export")
async def api_export(request: Request, background_tasks: BackgroundTasks):
    """Burn video + TTS audio + soft subtitle captions into one MP4.

    FIX: tmpdir is cleaned up via BackgroundTasks AFTER the response is
    fully streamed — prevents the race condition where FileResponse tries
    to read a file that Python GC already deleted.
    """
    import wave as _wave

    t0 = time.time()
    log.info("[export] Request received")
    tmpdir = None

    try:
        body  = await request.json()
        sid   = body.get("session_id", "")
        vibe  = body.get("vibe", "")
        lines = body.get("lines", [])

        log.info(f"[export] sid={sid} vibe={vibe} lines={len(lines)}")

        sess = SESSIONS.get(sid)
        if not sess:
            return JSONResponse({"ok": False, "error": "session not found"}, status_code=404)
        if not lines:
            return JSONResponse({"ok": False, "error": "no lines"}, status_code=400)

        video_path = sess["norm"]
        tmpdir     = tempfile.mkdtemp(prefix="mq_export_")
        log.info(f"[export] Working dir: {tmpdir}")

        # 1. Write TTS WAV files
        log.info("[export] Writing TTS WAV files…")
        audio_clips = []
        for i, line in enumerate(lines):
            b64 = line.get("audio_b64")
            if not b64:
                continue
            wav_bytes = base64.b64decode(b64)
            wav_path  = os.path.join(tmpdir, f"tts_{i:03d}.wav")
            with open(wav_path, "wb") as f:
                f.write(wav_bytes)
            with _wave.open(io.BytesIO(wav_bytes)) as wf:
                dur = wf.getnframes() / wf.getframerate()
            audio_clips.append({"path": wav_path, "t": float(line["t"]), "duration": dur})
        log.info(f"[export] {len(audio_clips)} audio clips written")

        # 2. Build SRT subtitle file
        log.info("[export] Building subtitle file…")
        total_dur  = sess["meta"]["n_frames"] / sess["meta"]["fps"]
        vibe_label = {
            "football": "COCO", "diva": "COCO", "wildlife": "COCO",
            "boxing": "COCO", "masterchef": "COCO",
        }.get(vibe, "COCO")

        def _srt_time(s):
            h=int(s//3600); m=int((s%3600)//60); sec=s%60
            ms=int((sec%1)*1000); sec=int(sec)
            return f"{h:02d}:{m:02d}:{sec:02d},{ms:03d}"

        srt_rows = []
        for i, line in enumerate(lines):
            t_start = float(line["t"])
            t_end   = float(lines[i+1]["t"]) if i+1<len(lines) else min(t_start+4, total_dur)
            srt_rows += [
                str(i+1),
                f"{_srt_time(t_start)} --> {_srt_time(t_end)}",
                f"{vibe_label}  {line['text']}",
                "",
            ]
        srt_path = os.path.join(tmpdir, "subs.srt")
        with open(srt_path, "w", encoding="utf-8") as f:
            f.write("\n".join(srt_rows))
        log.info(f"[export] SRT written: {len(lines)} cues")

        # 3. Build ffmpeg command
        out_path = os.path.join(tmpdir, "broadcast.mp4")
        n        = len(audio_clips)

        cmd = ["ffmpeg", "-y", "-i", video_path, "-i", srt_path]
        for clip in audio_clips:
            cmd += ["-i", clip["path"]]

        if n > 0:
            parts, delayed = [], []
            for idx, clip in enumerate(audio_clips):
                delay_ms = int(clip["t"] * 1000)
                lbl = f"d{idx}"
                cmd_idx = idx + 2   # 0=video, 1=srt, 2..=wavs
                parts.append(f"[{cmd_idx}:a]adelay={delay_ms}|{delay_ms},apad[{lbl}]")
                delayed.append(f"[{lbl}]")
            parts.append(f"{''.join(delayed)}amix=inputs={n}:normalize=0[commentary]")
            parts.append("[0:a]volume=0.25[orig];[orig][commentary]amix=inputs=2:normalize=0[mixed]")
            cmd += [
                "-filter_complex", ";".join(parts),
                "-map", "0:v",
                "-map", "[mixed]",
                "-map", "1:s",
                "-c:v", "copy",         # stream-copy video — no re-encode
                "-c:a", "aac", "-b:a", "128k",
                "-c:s", "mov_text",
                "-movflags", "+faststart",
                out_path,
            ]
        else:
            cmd += [
                "-map", "0:v", "-map", "0:a", "-map", "1:s",
                "-c:v", "copy", "-c:a", "copy",
                "-c:s", "mov_text",
                "-movflags", "+faststart",
                out_path,
            ]

        log.info(f"[export] Running ffmpeg (n_audio={n})…")
        log.info(f"[export] CMD: {' '.join(cmd)}")
        t1  = time.time()
        res = subprocess.run(cmd, capture_output=True)
        log.info(f"[export] ffmpeg finished in {time.time()-t1:.1f}s, "
                 f"returncode={res.returncode}")

        if res.returncode != 0:
            err = res.stderr.decode(errors="replace")[-1200:]
            log.error(f"[export] ffmpeg STDERR:\n{err}")
            return JSONResponse({"ok": False, "error": f"ffmpeg: {err}"}, status_code=500)

        out_size_mb = os.path.getsize(out_path) / 1024 / 1024
        log.info(f"[export] Output: {out_path} ({out_size_mb:.1f}MB)")
        log.info(f"[export] Total time: {time.time()-t0:.1f}s")

        vibe_slug = vibe or "broadcast"
        fname     = f"marquee_{vibe_slug}.mp4"

        # KEY FIX: schedule tmpdir cleanup as a background task so it runs
        # AFTER the file has been fully streamed to the client.
        # Previously tmpdir could be cleaned up mid-stream causing download to hang.
        def _cleanup(d: str):
            try:
                shutil.rmtree(d, ignore_errors=True)
                log.info(f"[export] Cleaned up tmpdir: {d}")
            except Exception as ex:
                log.warning(f"[export] Cleanup failed: {ex}")

        background_tasks.add_task(_cleanup, tmpdir)

        return FileResponse(
            out_path,
            media_type="video/mp4",
            filename=fname,
            headers={"Content-Disposition": f'attachment; filename="{fname}"'},
        )

    except Exception as e:
        import traceback; traceback.print_exc()
        log.error(f"[export] FAILED: {e}")
        if tmpdir:
            shutil.rmtree(tmpdir, ignore_errors=True)
        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)


@app.api(name="generate")
def generate(session_id: str, names: list, vibe: str) -> dict:
    t0 = time.time()
    log.info(f"[generate] session={session_id} vibe={vibe} names={names}")

    sess = SESSIONS.get(session_id)
    if not sess:
        log.error(f"[generate] Session not found: {session_id}")
        return {"ok": False, "error": "Session expired — re-cast"}

    tracks     = sess.get("tracks", [])
    id_to_name = {}
    for i, trk in enumerate(tracks):
        raw = str(names[i]).strip() if i < len(names) else ""
        if raw:
            id_to_name[trk["id"]] = raw

    log.info(f"[generate] Named tracks: {id_to_name}")
    if id_to_name:
        write_names_to_jsonl(sess["jsonl_path"], id_to_name)

    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)
    log.info(f"[generate] Key events: {event_times}")

    log.info("[generate] Annotating keyframes…")
    frames = []
    for ts, fidx in zip(event_times, event_idxs):
        bgr = frame_at(sess["norm"], fidx)
        if bgr is None:
            log.warning(f"[generate] frame_at({fidx}) returned None, skipping")
            continue
        frames.append((ts, annotate_event_frame(bgr, sess["jsonl_path"], ts)))
    log.info(f"[generate] {len(frames)} annotated frames ready for VLM")

    log.info(f"[generate] Calling Qwen VLM (style={style})…")
    t1     = time.time()
    roster = list(id_to_name.values())
    script = generate_commentary(frames, roster, style, vibe=vibe)
    log.info(f"[generate] VLM done in {time.time()-t1:.1f}s — {len(script)} lines")
    log.info(f"[generate] Total: {time.time()-t0:.1f}s")

    return {"ok": True,
            "script": [{"t": round(float(l["time"]),2), "text": l["text"]}
                       for l in script]}


app.launch(ssr_mode=False, show_error=True)