#!/usr/bin/env python3
"""Tighten a spoken take by removing dead air between phrases.

Talking-head takes carry pauses at line breaks — especially when read from a
teleprompter. This cuts them out so the VO carries the reel at pace.

Two engines, tried in order:
  auto-editor  (preferred)  pip install auto-editor — audio-level analysis
  ffmpeg       (fallback)   silencedetect + select/aselect, no extra install

Both preserve pitch: sections are cut, never sped up. Output goes to
work/extracts/ by convention; re-probe it before planning against it, because
its timeline no longer matches the original take.

DETERMINISM ACROSS RESOLUTIONS. auto-editor is not resolution-independent —
tightening the same audio in a 1080 and a 4K copy of one take can remove
different amounts, which breaks the "edit fast at low res, reproduce at 4K"
workflow because every face in-point then shifts. Capture the decisions once
and replay them instead:

  tighten_vo.py take_1080.mp4 --out t.mp4 --emit-cuts cuts.json
  tighten_vo.py take_4k.mp4  --out t4k.mp4 --apply-cuts cuts.json

--emit-cuts forces the ffmpeg analyser, which reads audio only and so gives
the same ranges at any resolution. --apply-cuts does no detection at all: it
replays the recorded ranges exactly.

Usage:
  tighten_vo.py raw/take_01.mp4 --out work/extracts/take_01_tight.mp4
  tighten_vo.py raw/take_01.mp4 --out work/extracts/t.mp4 --threshold 0.03 --margin 0.25
"""
import argparse, json, math, re, shutil, subprocess, sys
from pathlib import Path

# auto-editor's own defaults, restated so the ffmpeg fallback matches them.
DEFAULT_THRESHOLD = 0.04    # linear amplitude, 0-1
DEFAULT_MARGIN = 0.35       # seconds kept either side of speech.
                            # Below ~0.3 the cut eats word edges:
                            # a 0.25 margin clipped 'train' in the field.
# A pause is only dead air once it is longer than a thought boundary. Below
# this it is speech rhythm, and cutting it makes delivery sound rushed. 0.3
# was too eager: it removed the beat speakers naturally take mid-sentence.
MIN_SILENCE = 0.5
# A held pause immediately before a payoff line is doing work — mark those
# with --keep-pause-before and they survive even past MIN_SILENCE.
KEEP_WINDOW = 1.2           # how far before a marked time a pause is protected


def probe_duration(path):
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries",
                        "format=duration", "-of", "csv=p=0", str(path)],
                       capture_output=True, text=True)
    if r.returncode != 0 or not r.stdout.strip():
        sys.exit(f"Cannot read {path} — is it a media file ffmpeg can open?")
    return float(r.stdout.strip())


def run_auto_editor(src, out, threshold, margin):
    exe = shutil.which("auto-editor")
    if not exe:
        return False
    cmd = [exe, str(src), "-o", str(out), "--no-open",
           "--edit", f"audio:threshold={threshold}", "--margin", f"{margin}s"]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        print(f"auto-editor failed, falling back to ffmpeg:\n"
              f"{r.stderr.strip()[-500:]}", file=sys.stderr)
        return False
    return out.exists()


def detect_silences(src, threshold, min_silence):
    """Parse ffmpeg silencedetect output into (start, end) pairs.

    silencedetect logs at *info* level, so `-v error` would hide the very
    output we need — read stderr with the banner suppressed instead.
    """
    # silencedetect takes dB; convert the linear threshold auto-editor uses.
    noise_db = 20 * math.log10(threshold) if threshold > 0 else -60
    r = subprocess.run(
        ["ffmpeg", "-hide_banner", "-i", str(src), "-af",
         f"silencedetect=noise={noise_db:.1f}dB:d={min_silence}", "-f", "null", "-"],
        capture_output=True, text=True)
    starts = [float(m) for m in re.findall(r"silence_start: ([\d.]+)", r.stderr)]
    ends = [float(m) for m in re.findall(r"silence_end: ([\d.]+)", r.stderr)]
    return list(zip(starts, ends + [None] * (len(starts) - len(ends))))


def protected(silence, keep_before, window=KEEP_WINDOW):
    """Is this pause the beat before a payoff the operator marked?"""
    start, end = silence
    stop = end if end is not None else start
    return any(0 <= t - stop <= window for t in keep_before)


def keep_ranges(silences, duration, margin, keep_before=()):
    """Invert silence spans into the spans worth keeping, padded by margin.

    Pauses marked via keep_before are left intact: a held beat before a
    punchline is timing, not dead air, and removing it flattens the delivery.
    """
    silences = [s for s in silences if not protected(s, keep_before)]
    keeps, cursor = [], 0.0
    for start, end in silences:
        stop = min(start + margin, duration)
        if stop > cursor:
            keeps.append((cursor, stop))
        cursor = max(cursor, (end - margin) if end is not None else duration)
    if cursor < duration:
        keeps.append((cursor, duration))
    return [(a, b) for a, b in keeps if b - a > 0.01]


CUTS_SCHEMA = 1


def apply_keeps(src, out, keeps):
    """Cut the source down to `keeps` — the only place ranges become video."""
    expr = "+".join(f"between(t,{a:.3f},{b:.3f})" for a, b in keeps)
    tmp = out.with_name(out.stem + ".tmp" + out.suffix)
    r = subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(src),
         "-vf", f"select='{expr}',setpts=N/FRAME_RATE/TB",
         "-af", f"aselect='{expr}',asetpts=N/SR/TB",
         "-c:v", "libx264", "-preset", "medium", "-crf", "16",
         "-c:a", "aac", "-b:a", "192k", "-ar", "48000",
         "-movflags", "+faststart", str(tmp)],
        capture_output=True, text=True)
    if r.returncode != 0:
        tmp.unlink(missing_ok=True)
        sys.exit(f"ffmpeg tighten failed:\n{r.stderr[-800:]}")
    tmp.rename(out)
    return len(keeps)


def load_cuts(path, duration):
    p = Path(path)
    if not p.exists():
        sys.exit(f"Cut list not found: {p}. Produce one first with "
                 f"--emit-cuts on the take you edited.")
    try:
        data = json.loads(p.read_text())
    except json.JSONDecodeError as e:
        sys.exit(f"{p} is not valid JSON ({e}) — regenerate it with --emit-cuts.")
    if data.get("schema_version") != CUTS_SCHEMA:
        sys.exit(f"{path} is schema {data.get('schema_version')}, expected "
                 f"{CUTS_SCHEMA} — regenerate it with --emit-cuts.")
    keeps = [tuple(k) for k in data.get("keep", [])]
    if not keeps:
        sys.exit(f"{path} records no keep ranges — nothing to apply.")
    recorded = data.get("source_duration")
    if recorded and abs(recorded - duration) > 0.5:
        print(f"WARNING: cut list was captured from a {recorded:.2f}s source "
              f"but this one is {duration:.2f}s. Replaying anyway — if these "
              f"are different takes rather than two renders of one take, the "
              f"cuts will land in the wrong places.", file=sys.stderr)
    over = [k for k in keeps if k[1] > duration + 0.05]
    if over:
        sys.exit(f"{path} keeps audio past the end of this source "
                 f"({over[0][1]:.2f}s > {duration:.2f}s) — wrong take?")
    return keeps


def run_ffmpeg_fallback(src, out, threshold, margin, duration, keep_before=()):
    silences = detect_silences(src, threshold, MIN_SILENCE)
    keeps = keep_ranges(silences, duration, margin, keep_before)
    if not keeps:
        sys.exit("Every span read as silence — the take may be near-mute, or "
                 "--threshold is too high. Inspect the audio before retrying.")
    apply_keeps(src, out, keeps)
    return keeps


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("source")
    ap.add_argument("--out", required=True)
    ap.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD,
                    help=f"Loudness below this counts as silence, linear 0-1 "
                         f"(default {DEFAULT_THRESHOLD}). Raise it if breaths "
                         f"survive, lower it if words get clipped.")
    ap.add_argument("--margin", type=float, default=DEFAULT_MARGIN,
                    help=f"Seconds of silence kept either side of speech "
                         f"(default {DEFAULT_MARGIN}). Raise it if a word sounds clipped at the cut.")
    ap.add_argument("--engine", choices=("auto", "auto-editor", "ffmpeg"),
                    default="auto", help="Force an engine; default tries "
                                         "auto-editor then falls back.")
    ap.add_argument("--keep-pause-before", type=float, action="append",
                    default=[], metavar="SEC",
                    help="Protect the pause just before this time — the beat "
                         "before a payoff line is timing, not dead air. "
                         "Repeatable. Forces the ffmpeg analyser.")
    ap.add_argument("--strict", action="store_true",
                    help="Fail instead of substituting a fallback engine.")
    ap.add_argument("--emit-cuts", metavar="PATH",
                    help="Record the keep ranges to PATH. Forces the ffmpeg "
                         "analyser, which reads audio only and so produces "
                         "the same ranges at any resolution.")
    ap.add_argument("--apply-cuts", metavar="PATH",
                    help="Replay ranges from PATH instead of detecting. Use "
                         "this to reproduce a low-res edit exactly at 4K.")
    args = ap.parse_args()

    src = Path(args.source)
    if not src.exists():
        sys.exit(f"Source not found: {src}")
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    before = probe_duration(src)

    # Replay path: no detection at all, so the result is bit-identical in its
    # cut decisions to the run that produced the list.
    if args.apply_cuts:
        if args.emit_cuts:
            sys.exit("--emit-cuts and --apply-cuts are mutually exclusive.")
        keeps = load_cuts(args.apply_cuts, before)
        apply_keeps(src, out, keeps)
        after = probe_duration(out)
        print(json.dumps({
            "engine": "replay", "cuts_from": args.apply_cuts,
            "source": str(src), "output": str(out), "segments": len(keeps),
            "duration_before": round(before, 2),
            "duration_after": round(after, 2),
            "removed_seconds": round(before - after, 2),
        }, indent=2))
        return

    if args.emit_cuts and args.engine == "auto-editor":
        sys.exit("--emit-cuts needs the ffmpeg analyser to stay "
                 "resolution-independent; drop --engine auto-editor.")

    engine = None
    if args.engine in ("auto", "auto-editor") and not args.emit_cuts \
            and not args.keep_pause_before:
        if run_auto_editor(src, out, args.threshold, args.margin):
            engine = "auto-editor"
        elif args.engine == "auto-editor":
            sys.exit("auto-editor not available — install it "
                     "(pip install auto-editor) or use --engine ffmpeg")
    keeps = None
    if engine is None and args.strict:
        sys.exit("auto-editor unavailable and --strict was given. Install it "
                 "(pip install auto-editor) or drop --strict to use the "
                 "ffmpeg silencedetect path.")
    if engine is None:
        print("NOTE: auto-editor unavailable — using the ffmpeg silencedetect "
              "path. Measured equivalent on test material, but say so rather "
              "than reporting the result as auto-editor's.", file=sys.stderr)
        keeps = run_ffmpeg_fallback(src, out, args.threshold, args.margin,
                                    before, args.keep_pause_before)
        engine = "ffmpeg"

    if args.emit_cuts:
        dst = Path(args.emit_cuts)
        dst.parent.mkdir(parents=True, exist_ok=True)
        dst.write_text(json.dumps({
            "schema_version": CUTS_SCHEMA, "source": str(src),
            "source_duration": round(before, 3),
            "threshold": args.threshold, "margin": args.margin,
            "engine": engine,
            "keep": [[round(a, 3), round(b, 3)] for a, b in keeps],
        }, indent=2))
        print(f"Recorded {len(keeps)} keep ranges -> {dst}. Replay at any "
              f"resolution with --apply-cuts.", file=sys.stderr)

    after = probe_duration(out)
    summary = {
        "engine": engine, "engine_requested": args.engine,
        "substituted": args.engine == "auto" and engine != "auto-editor",
        "source": str(src), "output": str(out),
        "duration_before": round(before, 2), "duration_after": round(after, 2),
        "removed_seconds": round(before - after, 2),
        "removed_pct": round((before - after) / before * 100, 1) if before else 0,
    }
    print(json.dumps(summary, indent=2))
    print(f"\nRe-probe {out} before planning against it — its timeline no "
          f"longer matches {src.name}.", file=sys.stderr)


if __name__ == "__main__":
    main()