pipeline/assemble.py

"""Native audio-aware assembly (SPEC_NEW §7.4, §11.4).

Ported from the donor's `skills/assemble.md` Claude-code-execution skill into a
deterministic local module (numpy + stdlib wave + system ffmpeg). No LLM, no sandbox.

Fixed order (a correctness invariant, SPEC_NEW §12.7):
  1. per-clip trim to clean speech boundaries (peak-relative RMS + tail-junk cut)
  2. concat the trimmed clips (codec-identical → safe stream copy)
  3. two-pass loudnorm to LOUDNORM_TARGET_I (−14 LUFS)
  4. forced-aligned burned Swedish captions (best-effort; from kb-whisper word
     timestamps on the FINAL normalized audio → no caption drift)
  5. output media/final/{video_id}.mp4 (H.264 + AAC, 9:16, ≤ 60s)

`speech_window()` is the pure DSP core and is unit-tested directly with synthetic
signals (tests/test_assemble_stitch.py).
"""

from __future__ import annotations

import json
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from . import db, storage
from .config import MEDIA_DIR, get_settings

# ── Trim tuning (donor values, SPEC_NEW §11.4) ───────────────────────────────
TARGET_SR = 16000
FRAME_MS = 20
PEAK_FRAC = 0.15            # voiced iff rms > 0.15 * peak
GAP_CLOSE_MS = 250         # merge between-word pauses shorter than this
ISLAND_MS = 100            # drop isolated voiced blips shorter than this
TAIL_GAP_S = 0.60          # a trailing segment after a gap > this may be junk…
TAIL_KEEP_S = 1.5          # …unless the trailing voiced run reaches this length
PAD_START_S = 0.04
PAD_END_S = 0.18          # trailing silence kept after speech → a small gap before the next
                          # clip's speech, so cuts aren't too tight (owner: +~0.1s margin)
SANITY_FLOOR_S = 1.0       # if the speech window is shorter, keep the whole clip

TARGET_W, TARGET_H = 720, 1280   # 9:16 @ 720p — every clip normalized to this so the
TARGET_FPS = 30                  # concat is a safe stream copy and B-roll overlays line up

FINAL_DIR = MEDIA_DIR / "final"
_COVER_916 = (f"scale={TARGET_W}:{TARGET_H}:force_original_aspect_ratio=increase,"
              f"crop={TARGET_W}:{TARGET_H},fps={TARGET_FPS},setsar=1")


# ── Pure DSP core ─────────────────────────────────────────────────────────────

@dataclass
class Window:
    start: float
    end: float
    diag: dict[str, Any]


def _runs(mask: list[bool]) -> list[tuple[int, int, bool]]:
    """Collapse a boolean mask into (start_idx, end_idx_exclusive, value) runs."""
    runs: list[tuple[int, int, bool]] = []
    if not mask:
        return runs
    start = 0
    for i in range(1, len(mask)):
        if mask[i] != mask[start]:
            runs.append((start, i, mask[start]))
            start = i
    runs.append((start, len(mask), mask[start]))
    return runs


def speech_window(samples, sr: int = TARGET_SR, is_first: bool = False) -> Window:
    """Find [start, end] seconds of real speech in a clip's mono audio.

    Peak-relative voiced detection + gap-close + island-removal + gap-based
    tail-junk trimming (SPEC_NEW §11.4). Pure: no I/O. `samples` is a 1-D numpy
    array (any scale — the threshold is peak-relative, so scale-invariant)."""
    import numpy as np

    x = np.asarray(samples, dtype=np.float64).reshape(-1)
    hop = max(1, int(sr * FRAME_MS / 1000))
    dur = len(x) / sr if sr else 0.0
    nframes = len(x) // hop

    if nframes == 0:
        return Window(0.0, dur, {"empty": True, "n_segments": 0, "tail_dropped": 0})

    rms = np.empty(nframes)
    for i in range(nframes):
        frame = x[i * hop:(i + 1) * hop]
        rms[i] = float(np.sqrt(np.mean(frame * frame))) if frame.size else 0.0

    peak = float(rms.max())
    thr = peak * PEAK_FRAC
    voiced = [bool(v > thr) for v in rms] if peak > 0 else [False] * nframes
    voiced_frames = sum(voiced)

    fps = sr / hop  # frames per second
    gap_close = max(1, round(GAP_CLOSE_MS / FRAME_MS))
    island = max(1, round(ISLAND_MS / FRAME_MS))

    # Close short non-voiced gaps that are flanked by voiced frames.
    for (a, b, val) in _runs(voiced):
        if not val and a > 0 and b < nframes and (b - a) < gap_close:
            for i in range(a, b):
                voiced[i] = True
    # Remove short voiced islands.
    for (a, b, val) in _runs(voiced):
        if val and (b - a) < island:
            for i in range(a, b):
                voiced[i] = False

    segs = [(a / fps, b / fps) for (a, b, val) in _runs(voiced) if val]
    diag: dict[str, Any] = {
        "peak": round(peak, 4), "thr": round(thr, 4),
        "voiced_frames": voiced_frames, "total_frames": nframes,
        "n_segments": len(segs), "tail_dropped": 0, "gap_at_cut": 0.0,
    }

    if not segs:  # no detectable speech — keep the whole clip
        return Window(0.0, dur, {**diag, "fallback": "no_voiced"})

    # Gap-based tail-junk trim (walk backward).
    keep_until = len(segs) - 1
    trailing_voiced = 0.0
    dropped = 0
    gap_at_cut = 0.0
    for k in range(len(segs) - 1, 0, -1):
        gap = segs[k][0] - segs[k - 1][1]
        trailing_voiced += segs[k][1] - segs[k][0]
        if gap > TAIL_GAP_S and trailing_voiced < TAIL_KEEP_S:
            keep_until = k - 1
            dropped += 1
            gap_at_cut = gap
        else:
            break
    diag["tail_dropped"] = dropped
    diag["gap_at_cut"] = round(gap_at_cut, 2)

    start = segs[0][0]
    end = segs[keep_until][1]

    # Pads / clamps / sanity floor / first-clip force.
    start -= PAD_START_S
    end += PAD_END_S
    start = max(0.0, start)
    end = min(dur, end)
    if end - start < SANITY_FLOOR_S:
        diag["fallback"] = "below_floor"
        start, end = 0.0, dur
    if is_first:
        start = 0.0
    return Window(start, end, diag)


# ── ffmpeg I/O ────────────────────────────────────────────────────────────────

def _run(cmd: list[str], cwd: Path | None = None) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=True,
                          capture_output=True, text=True)


def _ffprobe_duration(path: str | Path) -> float:
    out = _run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
                "-of", "default=nw=1:nk=1", str(path)])
    return float(out.stdout.strip())


def _extract_wav(clip: str, wav: str, cwd: Path) -> None:
    # Force 16-bit PCM so the stdlib `wave` reader below has a known sample width.
    _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-i", clip,
          "-ac", "1", "-ar", str(TARGET_SR), "-vn", "-c:a", "pcm_s16le", "-f", "wav", wav], cwd=cwd)


def _load_wav(path: Path):
    """Read a mono PCM wav with the stdlib (no scipy). Returns (float array, sr).
    The trim threshold is peak-relative, so the normalization here is cosmetic."""
    import wave

    import numpy as np

    with wave.open(str(path), "rb") as wf:
        sr = wf.getframerate()
        n = wf.getnframes()
        width = wf.getsampwidth()
        nchan = wf.getnchannels()
        raw = wf.readframes(n)
    dtype = {1: np.int8, 2: np.int16, 4: np.int32}.get(width, np.int16)
    data = np.frombuffer(raw, dtype=dtype).astype(np.float32)
    if nchan > 1:
        data = data.reshape(-1, nchan).mean(axis=1)
    peak = float(np.max(np.abs(data))) or 1.0
    return data / peak, sr


def _trim(clip: str, out: str, start: float, length: float, cwd: Path) -> None:
    _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-ss", f"{start:.3f}",
          "-i", clip, "-t", f"{length:.3f}", "-vf", _COVER_916,
          "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
          "-c:a", "aac", "-ar", "48000", "-movflags", "+faststart", out], cwd=cwd)


def _overlay_broll(talking: str, broll: str, out: str, cwd: Path) -> None:
    """B-roll cutaway (SPEC_NEW §15.6): VIDEO = the B-roll clip (scaled-to-cover
    720×1280, looped/trimmed to the talking length), AUDIO = the talking clip's audio
    (the voiceover keeps playing). The B-roll's own audio is discarded."""
    dur = _ffprobe_duration(cwd / talking)
    _run(["ffmpeg", "-y", "-hide_banner", "-nostats",
          "-i", talking, "-stream_loop", "-1", "-i", broll,
          "-filter_complex", f"[1:v]{_COVER_916},format=yuv420p[bv]",
          "-map", "[bv]", "-map", "0:a", "-t", f"{dur:.3f}",
          "-c:v", "libx264", "-preset", "veryfast", "-pix_fmt", "yuv420p",
          "-c:a", "aac", "-ar", "48000", "-movflags", "+faststart", out], cwd=cwd)


def _broll_file_for(seg: dict[str, Any], idx: int, cwd: Path) -> str | None:
    """Local filename of the B-roll clip to overlay on this segment, or None."""
    bid = seg.get("broll_id")
    if not bid:
        return None
    row = db.fetch_one("select file_path from broll where id = %s and status = 'ready'", (str(bid),))
    fp = row and row.get("file_path")
    if not fp:
        return None
    name = f"broll_{idx:03d}.mp4"
    if fp.startswith(("http://", "https://")):
        import httpx
        r = httpx.get(fp, timeout=120)
        r.raise_for_status()
        (cwd / name).write_bytes(r.content)
    else:
        shutil.copyfile(fp, cwd / name)
    return name


def _concat(trimmed: list[str], out: str, cwd: Path) -> None:
    listfile = cwd / "list.txt"
    listfile.write_text("".join(f"file '{name}'\n" for name in trimmed))
    _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-f", "concat", "-safe", "0",
          "-i", "list.txt", "-c", "copy", "-movflags", "+faststart", out], cwd=cwd)


def _loudnorm_two_pass(src: str, out: str, target_i: float, cwd: Path) -> None:
    """Two-pass EBU R128 loudnorm to target_i LUFS (SPEC_NEW §7.4)."""
    flt = f"loudnorm=I={target_i}:TP=-1.5:LRA=11"
    pass1 = _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-i", src,
                  "-af", f"{flt}:print_format=json", "-f", "null", "-"], cwd=cwd)
    m = re.search(r"\{[^{}]*\"input_i\"[^{}]*\}", pass1.stderr, re.S)
    if not m:
        # Could not measure — fall back to single-pass; still far better than nothing.
        _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-i", src, "-af", flt,
              "-c:v", "copy", "-c:a", "aac", "-ar", "48000", out], cwd=cwd)
        return
    meas = json.loads(m.group(0))
    flt2 = (f"{flt}:measured_I={meas['input_i']}:measured_TP={meas['input_tp']}:"
            f"measured_LRA={meas['input_lra']}:measured_thresh={meas['input_thresh']}:"
            f"offset={meas['target_offset']}:linear=true")
    _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-i", src, "-af", flt2,
          "-c:v", "copy", "-c:a", "aac", "-ar", "48000", out], cwd=cwd)


# ── Captions (best-effort) ────────────────────────────────────────────────────

def _ass_ts(t: float) -> str:
    t = max(0.0, t)
    cs = int(round(t * 100))
    h, cs = divmod(cs, 360000)
    m, cs = divmod(cs, 6000)
    s, cs = divmod(cs, 100)
    return f"{h}:{m:02d}:{s:02d}.{cs:02d}"


def _build_ass(words: list[tuple[str, float, float]], max_per_line: int) -> str:
    header = (
        "[Script Info]\nScriptType: v4.00+\nPlayResX: 1080\nPlayResY: 1920\n\n"
        "[V4+ Styles]\n"
        "Format: Name, Fontname, Fontsize, PrimaryColour, OutlineColour, BackColour, "
        "Bold, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\n"
        "Style: Cap,Arial,72,&H00FFFFFF,&H00000000,&H64000000,-1,1,4,1,2,80,80,260,1\n\n"
        "[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n"
    )
    lines: list[str] = []
    for i in range(0, len(words), max_per_line):
        group = words[i:i + max_per_line]
        start = group[0][1]
        end = group[-1][2]
        text = " ".join(w[0] for w in group).replace("\n", " ")
        lines.append(f"Dialogue: 0,{_ass_ts(start)},{_ass_ts(end)},Cap,,0,0,0,,{text}")
    return header + "\n".join(lines) + "\n"


def _maybe_burn_captions(normalized: str, out: str, cwd: Path) -> bool:
    """Burn forced-aligned captions from kb-whisper word timestamps on the FINAL
    audio. Best-effort: any failure logs and leaves the caller to use `normalized`."""
    s = get_settings()
    if not s.captions_enabled:
        return False
    try:
        from . import qc

        words = qc.word_timestamps(cwd / normalized)
        if not words:
            print("[assemble] no word timestamps; skipping captions")
            return False
        ass_name = "captions.ass"
        (cwd / ass_name).write_text(_build_ass(words, s.caption_max_words_per_line))
        _run(["ffmpeg", "-y", "-hide_banner", "-nostats", "-i", normalized,
              "-vf", f"ass={ass_name}", "-c:v", "libx264", "-preset", "veryfast",
              "-pix_fmt", "yuv420p", "-c:a", "copy", "-movflags", "+faststart", out], cwd=cwd)
        return True
    except Exception as exc:  # noqa: BLE001 — captions are a best-effort final layer
        print(f"[assemble] caption burn failed ({type(exc).__name__}: {exc}); shipping without captions")
        return False


# ── Orchestration ─────────────────────────────────────────────────────────────

def _ordered_clips(video_id: str) -> list[dict[str, Any]]:
    rows = db.fetch_all(
        """
        select * from segments
         where video_id = %s and file_path is not null
           and coalesce(status, '') <> 'superseded'
         order by idx, take
        """,
        (video_id,),
    )
    if not rows:
        raise RuntimeError(f"video {video_id} has no generated, non-superseded segments to assemble")
    # Exactly one segment per idx survives the seed pick; guard against duplicates.
    seen: dict[int, dict[str, Any]] = {}
    for r in rows:
        if r["idx"] not in seen:
            seen[r["idx"]] = r
    return [seen[i] for i in sorted(seen)]


def assemble(video_id: str) -> str:
    """Stitch a video's segments into media/final/{video_id}.mp4 and mark it ready."""
    clips = _ordered_clips(video_id)
    s = get_settings()
    FINAL_DIR.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as tmp:
        cwd = Path(tmp)
        trimmed: list[str] = []

        for i, seg in enumerate(clips):
            fp = seg["file_path"]
            local = f"clip_{i:03d}.mp4"
            if storage.is_remote(fp):
                import httpx
                r = httpx.get(fp, timeout=180)
                r.raise_for_status()
                (cwd / local).write_bytes(r.content)
            else:
                src = Path(fp)
                if not src.exists():
                    raise RuntimeError(f"segment {seg['id']} file missing: {src}")
                shutil.copyfile(src, cwd / local)

            wav = f"audio_{i:03d}.wav"
            _extract_wav(local, wav, cwd)
            samples, sr = _load_wav(cwd / wav)
            clip_dur = _ffprobe_duration(cwd / local)
            win = speech_window(samples, sr, is_first=(i == 0))
            win.end = min(win.end, clip_dur)
            length = max(0.05, win.end - win.start)

            d = win.diag
            print(f"clip_{i:03d} energy: peak={d.get('peak', 0):.4f} thr={d.get('thr', 0):.4f} "
                  f"voiced_frames={d.get('voiced_frames', 0)}/{d.get('total_frames', 0)}")
            print(f"clip_{i:03d} segments: {d.get('n_segments', 0)} found, "
                  f"tail_dropped={d.get('tail_dropped', 0)} (gap={d.get('gap_at_cut', 0.0):.2f})")
            print(f"clip_{i:03d} trim: dur={clip_dur:.2f} speech={win.start:.2f}..{win.end:.2f} "
                  f"out_dur={length:.2f}")

            out = f"trimmed_{i:03d}.mp4"
            _trim(local, out, win.start, length, cwd)
            broll = _broll_file_for(seg, i, cwd)
            if broll:
                overlaid = f"overlaid_{i:03d}.mp4"
                _overlay_broll(out, broll, overlaid, cwd)
                print(f"clip_{i:03d} broll: overlaid {broll} over talking audio")
                trimmed.append(overlaid)
            else:
                trimmed.append(out)

        _concat(trimmed, "concat.mp4", cwd)
        _loudnorm_two_pass("concat.mp4", "normalized.mp4", s.loudnorm_target_i, cwd)

        final_local = "final.mp4"
        if not _maybe_burn_captions("normalized.mp4", "captioned.mp4", cwd):
            shutil.copyfile(cwd / "normalized.mp4", cwd / final_local)
        else:
            shutil.copyfile(cwd / "captioned.mp4", cwd / final_local)

        final_dest = FINAL_DIR / f"{video_id}.mp4"
        shutil.copyfile(cwd / final_local, final_dest)
        final_duration = _ffprobe_duration(final_dest)

    final_path_db = str(final_dest)
    if storage.enabled() and not get_settings().dry_run:
        try:
            final_path_db = storage.upload_bytes("finals", f"{video_id}.mp4",
                                                 final_dest.read_bytes(), "video/mp4")
        except Exception as exc:  # noqa: BLE001 — keep local on failure
            print(f"[storage] final upload failed ({exc}); keeping local")

    db.execute(
        "update videos set final_path = %s, duration_s = %s, status = 'ready' where id = %s",
        (final_path_db, round(final_duration, 3), video_id),
    )
    print(f"video {video_id}: assembled {final_path_db} ({final_duration:.2f}s, {len(clips)} clips)")
    return final_path_db