"""Background transcription: faster-whisper over each settled VOD's voice
track, segments landed in the FTS index for search and captions.

Optional dependency — without faster-whisper the worker just reports itself
unavailable and sleeps. Device autodetect (cuda → cpu int8); model and the
1-based voice track come from config.json ("whisper_model", default small;
"voice_track", default 1 = the combined mix). The DB is the job queue: a
vods row with no done/skipped transcripts row is pending work, so crashes
retry naturally and new VODs queue themselves.
"""
import os
import subprocess
import threading
import time
import uuid

from clipshelf import problems
from clipshelf import config
from clipshelf import db
from clipshelf.config import CFG, JOBS
POLL_S = 60
SETTLE_S = 300

_model = None
_model_name = None
_lock = threading.Lock()
# what the worker is doing right now, for /api/activity: the first run
# downloads a model (hundreds of MB) with no output, which looks broken
state = {"stage": "idle", "detail": "", "vod_id": None, "progress": 0.0}


def _drain(segs, total_ms: int) -> list:
    """Consume the segment generator, reporting how far in we are.

    faster-whisper yields lazily and every segment carries its end time, so
    honest progress costs nothing — it just was not being reported, which
    made a twenty-minute transcription look identical to a hung one."""
    rows = []
    for s in segs:
        rows.append((int(s.start * 1000), int((s.end - s.start) * 1000), s.text))
        if total_ms:
            state["progress"] = min(0.999, (s.end * 1000) / total_ms)
    state["progress"] = 1.0
    return rows


def available() -> bool:
    try:
        import faster_whisper  # noqa: F401
        return True
    except ImportError:
        return False


_cpu_only = False


def _get_model(force_cpu: bool = False):
    global _model, _model_name, _cpu_only
    from faster_whisper import WhisperModel
    want = str(CFG.get("whisper_model") or "small")
    if force_cpu:
        _cpu_only = True
        _model = None
    if _model is None or _model_name != want:
        state.update(stage="loading model", detail=want)
        if not _cpu_only:
            try:
                _model = WhisperModel(want, device="cuda", compute_type="int8_float16")
            except Exception:
                _cpu_only = True
        if _cpu_only:
            _model = WhisperModel(want, device="cpu", compute_type="int8")
        _model_name = want
        state.update(stage="idle", detail="")
    return _model


def _extract_voice(path: str, out_wav: str) -> bool:
    """16k mono wav of the configured voice track; falls back to the first
    audio track when the configured index doesn't exist."""
    track = max(1, int(CFG.get("voice_track") or 1)) - 1
    for amap in (f"0:a:{track}", "0:a:0"):
        r = subprocess.run(
            ["ffmpeg", "-nostdin", "-hide_banner", "-loglevel", "error", "-y",
             "-i", path, "-map", amap, "-ar", "16000", "-ac", "1", out_wav],
            capture_output=True, timeout=3600)
        if r.returncode == 0 and os.path.exists(out_wav):
            return True
    return False


def _set_status(conn, vod_id: int, status: str, **extra):
    conn.execute(
        """INSERT INTO transcripts (vod_id, status, model, error, segments, updated_at)
           VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
           ON CONFLICT(vod_id) DO UPDATE SET status = excluded.status,
               model = excluded.model, error = excluded.error,
               segments = excluded.segments, updated_at = CURRENT_TIMESTAMP""",
        (vod_id, status, extra.get("model"), extra.get("error"),
         extra.get("segments", 0)))
    conn.commit()


def store_segments(conn, vod_id: int, segments: list):
    """Replace a VOD's searchable segments: [(t_ms, dur_ms, text), ...]."""
    conn.execute("DELETE FROM transcript_fts WHERE vod_id = ?", (vod_id,))
    for t_ms, dur_ms, text in segments:
        text = " ".join(str(text).split())
        if text:
            conn.execute(
                "INSERT INTO transcript_fts (text, vod_id, t_ms, dur_ms) VALUES (?, ?, ?, ?)",
                (text, vod_id, int(t_ms), int(dur_ms)))
    conn.commit()


def transcribe_vod(vod_id: int) -> dict:
    conn = db.connect()
    try:
        vod = conn.execute("SELECT * FROM vods WHERE id = ?", (vod_id,)).fetchone()
        if not vod:
            return {"error": "no such vod"}
        path = os.path.join(config.VOD_DIR, vod["filename"])
        state.update(stage="extracting audio", detail=vod["filename"], vod_id=vod_id,
                     progress=0.0)
        _set_status(conn, vod_id, "running", model=str(CFG.get("whisper_model") or "small"))
        wav = str(JOBS / f"tw-{uuid.uuid4().hex[:8]}.wav")
        try:
            if not _extract_voice(path, wav):
                _set_status(conn, vod_id, "error", error="audio extraction failed")
                return {"error": "audio extraction failed"}
            with _lock:
                try:
                    model = _get_model()
                    state.update(stage="transcribing", detail=vod["filename"])
                    segs, _info = model.transcribe(wav, vad_filter=True)
                    rows = _drain(segs, vod["duration_ms"] or 0)
                except Exception as e:
                    # CUDA can construct fine and still die at inference
                    # (missing cublas, OOM) — fall back to CPU once
                    if _cpu_only:
                        raise
                    print(f"[transcribe] GPU path failed ({e}) — retrying on CPU", flush=True)
                    model = _get_model(force_cpu=True)
                    segs, _info = model.transcribe(wav, vad_filter=True)
                    rows = _drain(segs, vod["duration_ms"] or 0)
        finally:
            try:
                os.unlink(wav)
            except OSError:
                pass
        state.update(stage="idle", detail="", vod_id=None, progress=0.0)
        store_segments(conn, vod_id, rows)
        _set_status(conn, vod_id, "done", model=_model_name, segments=len(rows))
        print(f"[{time.strftime('%X')}] [transcribe] vod {vod_id}: {len(rows)} segments",
              flush=True)
        return {"segments": len(rows)}
    except Exception as e:
        problems.record("transcription", e,
                        "check the voice track in settings, or try a smaller whisper model")
        try:
            _set_status(conn, vod_id, "error", error=str(e)[:500])
        except Exception:
            pass
        return {"error": str(e)}
    finally:
        conn.close()


def why_pending(conn) -> str:
    """A one-line reason the queue looks stuck, for the UI."""
    if not available():
        return "faster-whisper isn't installed"
    rows = _pick_pending(conn)
    if not rows:
        return ""
    for r in rows:
        path = os.path.join(config.VOD_DIR, r["filename"])
        if not os.path.exists(path):
            continue
        age = time.time() - os.path.getmtime(path)
        if age < SETTLE_S:
            return (f"waiting for the recording to finish "
                    f"({int((SETTLE_S - age) / 60) + 1} min)")
        return "queued"
    return ""


def _pick_pending(conn):
    return conn.execute(
        """SELECT v.id, v.filename FROM vods v
           LEFT JOIN transcripts t ON t.vod_id = v.id
           WHERE v.duration_ms > 0
             AND (t.vod_id IS NULL OR t.status IN ('pending', 'running'))
           ORDER BY v.start_epoch DESC LIMIT 5""").fetchall()


def loop():
    time.sleep(90)
    while True:
        try:
            if available():
                conn = db.connect()
                try:
                    rows = _pick_pending(conn)
                finally:
                    conn.close()
                for r in rows:
                    path = os.path.join(config.VOD_DIR, r["filename"])
                    if not os.path.exists(path):
                        continue
                    if time.time() - os.path.getmtime(path) < SETTLE_S:
                        continue    # still recording
                    transcribe_vod(r["id"])
                    break           # one per cycle — fresh vods stay first
        except Exception as e:
            problems.record("transcription", e)
        time.sleep(POLL_S)


def start():
    threading.Thread(target=loop, daemon=True).start()
