#!/usr/bin/env python3
"""Surgically remove words or phrases from a spoken take.

tighten_vo.py removes silence. This removes SPEECH — a filler word, a fluffed
clause, a sentence that did not survive the edit — and re-joins the take with a
short held breath so the cut does not sound clipped.

Word boundaries come from word-level Whisper timestamps, the same input
gen_captions.py takes:
  mlx_whisper take.wav --word-timestamps True --output-format json

Two ways to say what goes. Both are index-free at the call site, so a
re-transcribe does not silently shift what you meant:

  --remove "to a plan"     match consecutive words (case- and punctuation-
                           insensitive). Ambiguous matches are refused, not
                           guessed — disambiguate with --occurrence.
  --remove-range 12:15     explicit half-open word indices, from --list

Cuts are applied through the same keep-range machinery tighten_vo.py uses, so
--emit-cuts records them and they replay at any resolution with
tighten_vo.py --apply-cuts.

Usage:
  trim_words.py words.json --video raw/take.mp4 --list
  trim_words.py words.json --video raw/take.mp4 --out work/extracts/t.mp4 \\
      --remove "to a plan" --remove "um"
"""
import argparse, json, re, sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from gen_captions import load_words                    # noqa: E402
from tighten_vo import apply_keeps, probe_duration, CUTS_SCHEMA   # noqa: E402

DEFAULT_BREATH = 0.12      # seconds of held silence left at each join
MIN_KEEP = 0.02
SENTENCE_END = re.compile(r"[.!?…]$|[—–-]$")
GROUP_GAP = 0.35           # a pause this long also ends a group
RESTART_PREFIX = 4         # shared leading words that mark a restart
FILLERS = {"um", "uh", "erm", "hmm"}


def sentence_groups(words, gap=GROUP_GAP):
    """Split into sentence-ish groups on punctuation OR a pause.

    Punctuation alone is not enough: a speaker who fluffs a line usually
    trails off without any, and Whisper punctuates an abandoned attempt
    inconsistently. The pause that follows a false start is the reliable
    signal, so a gap ends a group too.
    """
    groups, cur = [], []
    for i, w in enumerate(words):
        cur.append(i)
        ends = bool(SENTENCE_END.search(w["text"].strip()))
        if not ends and i + 1 < len(words):
            ends = (words[i + 1]["start"] - w["end"]) >= gap
        if ends:
            groups.append(cur)
            cur = []
    if cur:
        groups.append(cur)
    return groups


def common_prefix_length(left, right):
    n = 0
    for a, b in zip(left, right):
        if a != b:
            break
        n += 1
    return n


def find_restarts(words, groups, min_prefix=RESTART_PREFIX):
    """Word-index ranges for false starts the speaker later re-attempted.

    When someone fluffs a line they usually back up and start the sentence
    again with the same opening. So if a later group repeats an earlier
    group's opening — or one is wholly a prefix of the other — the EARLIER
    one is the abandoned attempt and the later one is the keeper.

    Borrowed from yap-editor (MIT, Jens Heitmann), whose mark_flubs() does
    the same comparison over sentence groups.
    """
    toks = [[norm(words[i]["text"]) for i in g] for g in groups]
    drop = []
    for earlier in range(len(groups)):
        for later in range(earlier + 1, len(groups)):
            l, r = toks[earlier], toks[later]
            if not l or not r:
                continue
            one_is_prefix = l == r[:len(l)] or r == l[:len(r)]
            if common_prefix_length(l, r) >= min_prefix or one_is_prefix:
                drop.append((groups[earlier][0], groups[earlier][-1] + 1))
                break
    return drop


def find_fillers(words):
    """Standalone filler words — 'um', 'uh' on their own."""
    return [(i, i + 1) for i, w in enumerate(words)
            if norm(w["text"]) in FILLERS]


def norm(s):
    return re.sub(r"[^\w']+", "", s).lower()


def find_phrase(words, phrase, occurrence=None):
    """Return the (start, end) word indices of a phrase, half-open."""
    target = [norm(w) for w in phrase.split() if norm(w)]
    if not target:
        sys.exit(f"--remove {phrase!r} has no matchable words.")
    normed = [norm(w["text"]) for w in words]
    hits = [i for i in range(len(normed) - len(target) + 1)
            if normed[i:i + len(target)] == target]
    if not hits:
        sys.exit(f"{phrase!r} not found in the transcript. Run --list to see "
                 f"the words as transcribed — Whisper may have misheard it.")
    if len(hits) > 1 and occurrence is None:
        spots = ", ".join(f"#{n + 1} at {words[i]['start']:.2f}s"
                          for n, i in enumerate(hits))
        sys.exit(f"{phrase!r} appears {len(hits)} times ({spots}). Pick one "
                 f"with --occurrence N rather than letting me guess.")
    i = hits[(occurrence - 1) if occurrence else 0]
    if occurrence and not 1 <= occurrence <= len(hits):
        sys.exit(f"--occurrence {occurrence} but {phrase!r} appears "
                 f"{len(hits)} time(s).")
    return i, i + len(target)


def cut_spans(words, removals, breath):
    """Turn word-index ranges into the source-time spans to drop."""
    spans = []
    for a, b in removals:
        start = words[a]["start"]
        end = words[b - 1]["end"]
        # Leave a breath on the outgoing side so the join is not abrupt; the
        # gap before the next word already supplies the incoming side.
        spans.append((max(0.0, start), max(0.0, end - breath)))
    spans.sort()
    merged = []
    for s in spans:
        if merged and s[0] <= merged[-1][1]:
            merged[-1] = (merged[-1][0], max(merged[-1][1], s[1]))
        else:
            merged.append(list(s) if False else (s[0], s[1]))
    return [(a, b) for a, b in merged if b > a]


def keeps_from_cuts(cuts, duration):
    keeps, cursor = [], 0.0
    for a, b in cuts:
        if a > cursor:
            keeps.append((cursor, a))
        cursor = max(cursor, b)
    if cursor < duration:
        keeps.append((cursor, duration))
    return [(a, b) for a, b in keeps if b - a > MIN_KEEP]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("words_json")
    ap.add_argument("--video", required=True)
    ap.add_argument("--out")
    ap.add_argument("--remove", action="append", default=[], metavar="PHRASE")
    ap.add_argument("--remove-range", action="append", default=[],
                    metavar="I:J", help="Half-open word indices from --list")
    ap.add_argument("--drop-restarts", action="store_true",
                    help="Drop sentences the speaker abandoned and re-attempted "
                         "— when a later sentence repeats an earlier one's "
                         "opening, the earlier one goes.")
    ap.add_argument("--drop-fillers", action="store_true",
                    help="Drop standalone filler words (um, uh, erm, hmm).")
    ap.add_argument("--dry-run", action="store_true",
                    help="Print what would be removed and stop. Always do this "
                         "first with --drop-restarts.")
    ap.add_argument("--occurrence", type=int,
                    help="Which match of --remove to take, 1-based")
    ap.add_argument("--breath", type=float, default=DEFAULT_BREATH,
                    help=f"Seconds of held silence left at each join "
                         f"(default {DEFAULT_BREATH}). Zero sounds clipped.")
    ap.add_argument("--emit-cuts", metavar="PATH",
                    help="Record keep ranges for tighten_vo.py --apply-cuts")
    ap.add_argument("--list", action="store_true",
                    help="Print indexed words and exit — no rendering")
    args = ap.parse_args()

    words = load_words(args.words_json)

    if args.list:
        for i, w in enumerate(words):
            print(f"{i:4d}  {w['start']:7.2f}-{w['end']:6.2f}  {w['text']}")
        return

    if not args.out and not args.dry_run:
        sys.exit("--out is required unless you passed --list or --dry-run.")
    auto = args.drop_restarts or args.drop_fillers
    if not args.remove and not args.remove_range and not auto:
        sys.exit("Nothing to remove — pass --remove PHRASE, --remove-range I:J, "
                 "--drop-restarts or --drop-fillers (--list shows the words).")

    removals = []
    for phrase in args.remove:
        removals.append(find_phrase(words, phrase, args.occurrence))
    for spec in args.remove_range:
        try:
            a, b = (int(n) for n in spec.split(":", 1))
        except ValueError:
            sys.exit(f"--remove-range {spec!r} must look like 12:15")
        if not 0 <= a < b <= len(words):
            sys.exit(f"--remove-range {spec!r} is outside 0:{len(words)}")
        removals.append((a, b))

    if args.drop_restarts:
        found = find_restarts(words, sentence_groups(words))
        for a, b in found:
            print(f"  restart: words {a}-{b - 1}  "
                  f"\"{' '.join(w['text'] for w in words[a:b])[:60]}\"")
        removals.extend(found)
    if args.drop_fillers:
        found = find_fillers(words)
        if found:
            print(f"  fillers: {len(found)} standalone "
                  f"({', '.join(words[a]['text'] for a, _ in found[:6])})")
        removals.extend(found)

    if args.dry_run:
        if not removals:
            print("Nothing matched — nothing would be removed.")
        return

    src = Path(args.video)
    if not src.exists():
        sys.exit(f"Video not found: {src}")
    duration = probe_duration(src)
    cuts = cut_spans(words, removals, args.breath)
    keeps = keeps_from_cuts(cuts, duration)
    if not keeps:
        sys.exit("Those removals would delete the whole take.")

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    apply_keeps(src, out, keeps)
    after = probe_duration(out)

    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(duration, 3),
            "threshold": None, "margin": args.breath, "engine": "trim_words",
            "keep": [[round(a, 3), round(b, 3)] for a, b in keeps],
        }, indent=2))

    print(json.dumps({
        "source": str(src), "output": str(out),
        "removed_phrases": args.remove, "removed_ranges": args.remove_range,
        "cuts": len(cuts), "segments": len(keeps),
        "duration_before": round(duration, 2),
        "duration_after": round(after, 2),
        "removed_seconds": round(duration - after, 2),
    }, indent=2))
    print(f"\nRe-transcribe or re-probe {out} before captioning against it — "
          f"every word timestamp after the first cut has moved.",
          file=sys.stderr)


if __name__ == "__main__":
    main()