pipeline/qc.py

"""Automatic QC per segment (SPEC_NEW §7.3), in order, on the RAW Veo clip:

1. Duration: within ±15% of the segment's OWN requested bucket
   (segments.duration_s ∈ {4,6,8}), not a flat 8 (SPEC_NEW §12.3). Fail → regen,
   with clip-clipping escalation (§12.1): if Veo honored the length but ASR is
   poor, the line was crammed → bump the segment up one bucket before regen.
2. ASR match: kb-whisper transcript vs spoken_text, both normalized, token-level
   Levenshtein ratio against thresholds from env.
3. Visual check: 4 sampled frames → Claude vision → JSON verdict. Any flag →
   qc_flag (never auto-fail in v1; we collect data first).

`qc_clip()` is the pure engine (used by golden tests); `run_qc()` wraps it with
DB load/store and the regen/attempt bookkeeping. `take_score()` ranks the two
hook takes for the seed pick (§11.3).
"""

from __future__ import annotations

import base64
import datetime as dt
import json
import subprocess
import tempfile
from pathlib import Path
from typing import Any

from psycopg.types.json import Json

from . import db, duration
from .config import get_settings
from .swedish import tokens

EXPECTED_SEGMENT_SECONDS = 8.0
DURATION_TOLERANCE = 0.15

VISUAL_CHECK_SCHEMA = {
    "type": "object",
    "properties": {
        "garbled_text_in_frame": {"type": "boolean"},
        "deformed_hands_or_face": {"type": "boolean"},
        "wrong_product": {"type": "boolean"},
        "watermark": {"type": "boolean"},
        "notes": {"type": "string"},
    },
    "required": [
        "garbled_text_in_frame",
        "deformed_hands_or_face",
        "wrong_product",
        "watermark",
        "notes",
    ],
    "additionalProperties": False,
}


# ── Duration ─────────────────────────────────────────────────────────────────

def probe_duration(path: str | Path) -> float:
    out = subprocess.run(
        [
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            str(path),
        ],
        capture_output=True, text=True, check=True,
    )
    return float(out.stdout.strip())


def duration_ok(actual_s: float, expected_s: float = EXPECTED_SEGMENT_SECONDS) -> bool:
    # epsilon keeps the boundary inclusive despite float representation
    return abs(actual_s - expected_s) <= DURATION_TOLERANCE * expected_s + 1e-9


# ── ASR (kb-whisper via faster-whisper) ──────────────────────────────────────

_whisper_model = None
_whisper_model_name: str | None = None


def _load_whisper():
    """Lazy singleton. Tries the configured (large) model, falls back to medium
    int8 on failure — per SPEC §3 stack table."""
    global _whisper_model, _whisper_model_name
    if _whisper_model is not None:
        return _whisper_model

    import os

    from faster_whisper import WhisperModel

    s = get_settings()
    device = s.kb_whisper_device
    if device == "auto":
        device = "cpu"  # faster-whisper/CTranslate2 has no MPS backend; cpu on macs
    compute = s.kb_whisper_compute_type or ("float16" if device == "cuda" else "int8")

    for model_name in (s.kb_whisper_model, s.kb_whisper_fallback_model):
        try:
            _whisper_model = WhisperModel(
                model_name,
                device=device,
                compute_type=compute,
                cpu_threads=min(8, os.cpu_count() or 4),
            )
            _whisper_model_name = model_name
            return _whisper_model
        except Exception as exc:  # noqa: BLE001 — fall through to the smaller model, then fail loud
            last_exc = exc
    raise RuntimeError(
        f"could not load any kb-whisper model "
        f"({s.kb_whisper_model}, fallback {s.kb_whisper_fallback_model}): {last_exc}"
    )


def transcribe(path: str | Path) -> str:
    model = _load_whisper()
    segments, _info = model.transcribe(
        str(path),
        language="sv",
        beam_size=5,
        vad_filter=True,
        vad_parameters={"min_silence_duration_ms": 500},
        condition_on_previous_text=False,  # KBLab's own anti-hallucination recommendation
    )
    return " ".join(seg.text.strip() for seg in segments).strip()


def word_timestamps(path: str | Path) -> list[tuple[str, float, float]]:
    """Per-word (text, start_s, end_s) for caption alignment (SPEC_NEW §11.4).

    Uses faster-whisper word timestamps — already a dependency — instead of
    pulling stable-ts/WhisperX. Swap in WhisperX here if higher precision is
    needed; the assembler only depends on this shape.
    """
    model = _load_whisper()
    segments, _info = model.transcribe(
        str(path), language="sv", beam_size=5, vad_filter=True,
        word_timestamps=True, condition_on_previous_text=False,
    )
    out: list[tuple[str, float, float]] = []
    for seg in segments:
        for w in (getattr(seg, "words", None) or []):
            text = (w.word or "").strip()
            if text:
                out.append((text, float(w.start), float(w.end)))
    return out


# ── Similarity ───────────────────────────────────────────────────────────────

def token_levenshtein_ratio(a: list[str], b: list[str]) -> float:
    """1 - levenshtein(a, b) / max(len). Token-level, per SPEC §7.3."""
    if not a and not b:
        return 1.0
    if not a or not b:
        return 0.0
    prev = list(range(len(b) + 1))
    for i, ta in enumerate(a, start=1):
        cur = [i] + [0] * len(b)
        for j, tb in enumerate(b, start=1):
            cur[j] = min(
                prev[j] + 1,          # deletion
                cur[j - 1] + 1,       # insertion
                prev[j - 1] + (ta != tb),  # substitution
            )
        prev = cur
    return 1.0 - prev[-1] / max(len(a), len(b))


def similarity(expected_text: str, transcript: str) -> float:
    return token_levenshtein_ratio(tokens(expected_text), tokens(transcript))


# ── Visual check (Claude vision) ─────────────────────────────────────────────

def sample_frames(path: str | Path, count: int = 4) -> list[bytes]:
    """Evenly spaced JPEG frames from the clip."""
    dur = probe_duration(path)
    frames: list[bytes] = []
    with tempfile.TemporaryDirectory() as tmp:
        for i in range(count):
            ts = dur * (i + 0.5) / count
            out = Path(tmp) / f"frame_{i}.jpg"
            subprocess.run(
                ["ffmpeg", "-v", "error", "-ss", f"{ts:.3f}", "-i", str(path),
                 "-frames:v", "1", "-q:v", "3", str(out)],
                check=True, capture_output=True,
            )
            frames.append(out.read_bytes())
    return frames


def visual_check(path: str | Path, brief_context: str = "") -> dict[str, Any]:
    """4 frames → Claude vision → JSON verdict. Respects DRY_RUN (SPEC §5)."""
    s = get_settings()
    if s.dry_run:
        return {"skipped": "dry_run"}

    import anthropic

    s.require("ANTHROPIC_API_KEY")
    client = anthropic.Anthropic(api_key=s.anthropic_api_key)

    content: list[dict[str, Any]] = []
    for frame in sample_frames(path):
        content.append({
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": "image/jpeg",
                "data": base64.standard_b64encode(frame).decode(),
            },
        })
    content.append({
        "type": "text",
        "text": (
            "These are 4 frames sampled evenly from an AI-generated vertical UGC ad clip "
            "(one person speaking Swedish to camera).\n"
            + (f"Creative brief context: {brief_context}\n" if brief_context else "")
            + "Inspect the frames for generation defects and answer strictly as JSON:\n"
            "- garbled_text_in_frame: any nonsense/garbled text visible in frame\n"
            "- deformed_hands_or_face: anatomically wrong hands, face, or teeth\n"
            "- wrong_product: a product is shown that contradicts the brief context "
            "(false if no brief context or no product visible)\n"
            "- watermark: any watermark or logo overlay\n"
            "- notes: one short sentence of justification"
        ),
    })

    response = client.messages.create(
        model=s.anthropic_model,
        max_tokens=1024,
        messages=[{"role": "user", "content": content}],
        output_config={"format": {"type": "json_schema", "schema": VISUAL_CHECK_SCHEMA}},
    )
    text = next((b.text for b in response.content if b.type == "text"), None)
    if text is None:
        return {"error": f"no text in visual verdict (stop_reason={response.stop_reason})"}
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return {"error": f"unparseable visual verdict: {text[:200]}"}


def visual_flagged(visual: dict[str, Any]) -> bool:
    if "skipped" in visual:
        return False
    if "error" in visual:
        return True
    return any(
        visual.get(k) is True
        for k in ("garbled_text_in_frame", "deformed_hands_or_face", "wrong_product", "watermark")
    )


# ── The QC engine ────────────────────────────────────────────────────────────

def qc_clip(
    path: str | Path,
    expected_text: str,
    expected_duration_s: float = EXPECTED_SEGMENT_SECONDS,
    brief_context: str = "",
    run_visual: bool = True,
) -> dict[str, Any]:
    """Run the full QC pipeline on one clip. Returns the qc JSON object."""
    s = get_settings()
    actual_s = probe_duration(path)
    dur_ok = duration_ok(actual_s, expected_duration_s)

    try:
        transcript = transcribe(path)
        sim: float | None = similarity(expected_text, transcript)
        asr_ok = True
    except Exception as exc:  # noqa: BLE001 — faster-whisper unavailable (lean deploy) → skip ASR
        transcript, sim, asr_ok = "", None, False
        print(f"[QC] ASR unavailable ({type(exc).__name__}); scoring on duration + visual only")

    visual = visual_check(path, brief_context) if run_visual else {"skipped": "disabled"}

    if asr_ok:
        if not dur_ok or sim < s.qc_flag_threshold:
            verdict = "fail"
        elif sim < s.qc_pass_threshold or visual_flagged(visual):
            verdict = "flag"
        else:
            verdict = "pass"
    else:
        # No ASR (lean deploy) — verdict from duration + visual only.
        verdict = "fail" if not dur_ok else ("flag" if visual_flagged(visual) else "pass")

    return {
        "transcript": transcript,
        "expected_text": expected_text,
        "similarity": round(sim, 4) if sim is not None else None,
        "duration_s": round(actual_s, 3),
        "expected_duration_s": expected_duration_s,
        "duration_ok": dur_ok,
        "visual": visual,
        "visual_ok": not visual_flagged(visual),
        "verdict": verdict,
        "asr_model": _whisper_model_name if asr_ok else None,
        "checked_at": dt.datetime.now(dt.timezone.utc).isoformat(),
    }


def take_score(qc: dict[str, Any] | None) -> float:
    """Rank a hook take for the seed pick (SPEC_NEW §11.3). Higher is better:
    primary = ASR similarity; small bonus for a clean visual; a non-passing
    verdict is strongly penalized so a failing take never wins."""
    if not qc:
        return -1.0
    sim = float(qc.get("similarity") or 0.0)
    vis_bonus = 0.005 if qc.get("visual_ok", True) else 0.0
    verdict_penalty = 0.0 if qc.get("verdict") in ("pass", "flag") else -1.0
    return sim + vis_bonus + verdict_penalty


def run_qc(segment_id: str) -> dict[str, Any]:
    """QC a segment from the DB, store the verdict, drive regen bookkeeping."""
    s = get_settings()
    seg = db.fetch_one("select * from segments where id = %s", (segment_id,))
    if seg is None:
        raise RuntimeError(f"no segment {segment_id}")
    if not seg["file_path"] or not Path(seg["file_path"]).exists():
        raise RuntimeError(f"segment {segment_id} has no downloaded file to QC")

    video = db.fetch_one("select * from videos where id = %s", (seg["video_id"],))
    brief_context = json.dumps(video["tags"], ensure_ascii=False) if video else ""

    # Per-segment expected duration (the requested bucket), NOT a flat 8 (§12.3).
    expected_dur = float(seg.get("duration_s") or EXPECTED_SEGMENT_SECONDS)
    result = qc_clip(
        seg["file_path"], seg["spoken_text"],
        expected_duration_s=expected_dur, brief_context=brief_context,
    )

    status = {"pass": "qc_pass", "flag": "qc_flag", "fail": "qc_fail"}[result["verdict"]]
    db.execute(
        "update segments set qc = %s, status = %s where id = %s",
        (Json(result), status, segment_id),
    )

    if result["verdict"] == "fail":
        if seg["attempts"] >= s.max_segment_attempts:
            db.execute("update videos set status = 'qc_failed' where id = %s", (seg["video_id"],))
        else:
            # Clip-clipping escalation (§12.1): Veo honored the length but speech
            # was crammed (good duration, poor ASR) → the line needs a longer clip.
            if result["duration_ok"] and result["similarity"] is not None and result["similarity"] < s.qc_flag_threshold:
                nb = duration.next_bucket(int(expected_dur))
                if nb is not None:
                    db.execute(
                        "update segments set duration_s = %s where id = %s", (nb, segment_id)
                    )
                    print(f"[QC] segment {segment_id}: crammed speech → bumped {int(expected_dur)}s → {nb}s")
            db.enqueue("generate_segment", {"segment_id": segment_id})
    return result