#!/usr/bin/env python3
"""Propose a beat-synced cut plan from normalized clips + beat data.

Pacing model — phase-based, all durations snapped to actual beat timestamps:

  HOOK       ~2 bars hold      grab attention; most kinetic clip first
  ESTABLISH  4-beat cuts       let shots read; set the scene
  GROOVE     2-beat cuts       settle into rhythm
  PUNCH      2-beat cuts       energy peak (1-beat only via --punch-beats 1 —
                               sub-second cuts read as "busy", so opt-in)
  BREATHE    4-beat cuts       relief after the peak
  BUILD      2-beat cuts       re-accelerate into the ending
  CLOSE      ~2 bars hold      payoff / smile / CTA

Guarantees:
  - minimum cut length is 2 beats (never the 1-beat flicker) unless --punch-beats 1
  - in-points are spread across each clip's FULL duration, not clustered at the start
  - two adjacent segments never come from the same clip (that's a jump cut)
  - clips are picked least-recently-used, not in strict rotation, so a long reel
    does not visibly cycle through the same order and reuse spreads out
  - in-points are positional by default. --score-in-points tries several
    candidates per segment and keeps the one that actually contains a person,
    using Apple Vision (see scripts/detect_subjects.py). On real footage that
    separates cleanly where edge density did not.

The output plan/cut_plan.json is meant to be HUMAN-EDITED before rendering.
Every field is safe to change: reorder segments, change "in" points, durations
(renderer re-snaps audio automatically), swap clip paths.

Usage:
  build_cut_plan.py work/clips.json work/beats.json --target-duration 30 \
      --mezz-dir work/mezz --out plan/cut_plan.json \
      [--order clipA,clipB,...] [--hook clipC] [--close clipD] [--punch-beats 1]
"""
import argparse, json, random, re, subprocess, sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

CANDIDATES = 4        # in-points tried per segment when scoring content
SRT_TIME = re.compile(r"(\d\d):(\d\d):(\d\d)[,.](\d{1,3})\s*-->\s*"
                      r"(\d\d):(\d\d):(\d\d)[,.](\d{1,3})")


def parse_srt(text):
    """[(start, end, line)] from an SRT. Tolerates comma or dot decimals."""
    cues, cur = [], None
    for raw in text.splitlines():
        line = raw.strip()
        m = SRT_TIME.search(line)
        if m:
            g = [int(x) for x in m.groups()]
            start = g[0] * 3600 + g[1] * 60 + g[2] + g[3] / 1000
            end = g[4] * 3600 + g[5] * 60 + g[6] + g[7] / 1000
            cur = [start, end, []]
            cues.append(cur)
        elif cur is not None and line and not line.isdigit():
            cur[2].append(line)
    out = []
    for s, e, parts in cues:
        # strip karaoke/position tags youtube auto-captions carry
        txt = re.sub(r"<[^>]+>", "", " ".join(parts)).strip()
        if txt:
            out.append((s, e, txt))
    return out


def find_lyric(cues, phrase, occurrence=None):
    """The cue whose text contains `phrase`. Ambiguity is refused, not guessed."""
    want = re.sub(r"[^\w ]+", "", phrase).lower().split()
    if not want:
        sys.exit(f"--music-start-lyric {phrase!r} has no searchable words.")
    hits = []
    for s, e, txt in cues:
        flat = re.sub(r"[^\w ]+", "", txt).lower()
        if " ".join(want) in flat:
            hits.append((s, e, txt))
    if not hits:
        sys.exit(f"{phrase!r} is not in the lyrics. Check the .srt — auto "
                 f"captions mishear a lot.")
    if len(hits) > 1 and occurrence is None:
        spots = ", ".join(f"#{i+1} at {s:.2f}s" for i, (s, _, _) in enumerate(hits))
        sys.exit(f"{phrase!r} appears {len(hits)} times ({spots}). Pick one "
                 f"with --lyric-occurrence N.")
    idx = (occurrence - 1) if occurrence else 0
    if not 0 <= idx < len(hits):
        sys.exit(f"--lyric-occurrence {occurrence} but {phrase!r} appears "
                 f"{len(hits)} time(s).")
    return hits[idx]


def snap_to_beat(t, beats, strong=None, prefer_strong=True):
    """Nearest beat to t. Prefers a downbeat when one is about as close.

    A lyric lands wherever the singer started; cutting there is off-grid and
    reads as sloppy. Snapping keeps the edit tight. Downbeats are the stronger
    cut point, so a downbeat wins unless a plain beat is clearly closer.
    """
    if not beats:
        return t
    nearest = min(beats, key=lambda b: abs(b - t))
    if prefer_strong and strong:
        s = min(strong, key=lambda b: abs(b - t))
        if abs(s - t) <= abs(nearest - t) + 0.12:
            return s
    return nearest
FLAT_SCORE = 0.01     # detect_subjects returns 0.0 when nobody is there


def phase_order(clips, phase_index):
    """A deterministic per-phase permutation of the clip pool.

    Least-recently-used is NOT enough on its own: with no-adjacent-repeat and
    even usage, LRU produces exactly the same cycle strict rotation does, which
    is the "reel visibly loops" complaint. Reshuffling the pool per phase makes
    each phase run a different order, so a 14-clip reel stops replaying one
    sequence four times. Seeded by phase index, so plans stay reproducible.
    """
    order = sorted(clips, key=lambda c: c["file"])
    rng = random.Random(1000 + phase_index)
    rng.shuffle(order)
    return order


def pick_clip(clips, prev, last_used, dur, phase_index):
    """Least-used clip, but scanned in this phase's own order.

    LRU spreads reuse; the per-phase order breaks the macro-pattern. Together
    they give even usage without an audible loop.
    """
    eligible = [c for c in phase_order(clips, phase_index)
                if (c is not prev or len(clips) == 1)
                and c["duration"] >= dur + 0.1]
    if not eligible:
        return None
    fewest = min(last_used.get(id(c), (0, -1))[0] for c in eligible)
    return next(c for c in eligible
                if last_used.get(id(c), (0, -1))[0] == fewest)


def candidate_times(clip_dur, seg_dur, usage_i, usage_n, n_candidates):
    """Start times to try for usage `usage_i` of `usage_n`, inside its window.

    The window keeps the spread-across-the-clip guarantee; the candidates
    inside it are what let content scoring avoid a dead frame.
    """
    span = max(0.0, clip_dur - seg_dur - 0.1)
    if span <= 0:
        return [0.0]
    width = span / max(1, usage_n)
    lo = usage_i * width
    if n_candidates <= 1:
        return [round(lo, 2)]
    return [round(min(span, lo + width * k / (n_candidates - 1)), 2)
            for k in range(n_candidates)]


def best_candidate(scores):
    """Index of the highest-scoring candidate; earliest wins ties.

    All-flat means every option is a dead frame — the caller should say so
    rather than silently pick one.
    """
    if not scores:
        return 0, False
    best = max(range(len(scores)), key=lambda i: (scores[i], -i))
    return best, scores[best] >= FLAT_SCORE


def content_score(path, t):
    """0.0 when no person is in frame, otherwise confidence, centre-weighted.

    This used to be edge density, which does not work: on real go-kart 4K an
    empty track scored 6.3-12.6 against 8.6-12.4 for helmet close-ups, because
    tarmac carries gravel, kerbs, fences and a horizon. Apple Vision separates
    the same frames with no overlap at all — 0 detections on every empty frame,
    0.65-0.75 confidence on every subject. See scripts/detect_subjects.py.
    """
    from detect_subjects import score_at
    return score_at(path, [t])[0][1]

# (name, beats_per_cut_multiplier, share of body time). Multiplier is applied
# to the minimum cut length: with the default 2-beat minimum, establish/breathe
# are 4-beat and groove/punch/build are 2-beat.
BODY_PHASES = [
    ("establish", 2, 0.22),
    ("groove",    1, 0.24),
    ("punch",     0, 0.16),   # 0 = punch_beats (the fastest allowed cut)
    ("breathe",   2, 0.16),
    ("build",     1, 0.22),
]


def resolve_music_start(args, beats):
    """(music_start, human note). Both flags snap; neither means 0.0."""
    times = beats.get("beats") or []
    strong = beats.get("strong_beats") or []
    if args.music_start_lyric:
        srt = args.lyrics
        if not srt:
            near = sorted(Path(beats["audio"]).parent.glob("*.srt"))
            if not near:
                sys.exit("No .srt beside the track — pass --lyrics, or fetch "
                         "the track with scripts/fetch_music.py which saves "
                         "subtitles when the video has them.")
            srt = str(near[0])
        cues = parse_srt(Path(srt).read_text(errors="replace"))
        if not cues:
            sys.exit(f"{srt} has no readable cues.")
        s, _, txt = find_lyric(cues, args.music_start_lyric,
                               args.lyric_occurrence)
        snapped = snap_to_beat(s, times, strong)
        return round(snapped, 3), (f'lyric "{txt[:40]}" at {s:.2f}s -> beat at '
                                   f'{snapped:.2f}s')
    if args.music_start is not None:
        snapped = snap_to_beat(args.music_start, times, strong)
        return round(snapped, 3), (f"{args.music_start:.2f}s -> beat at "
                                   f"{snapped:.2f}s")
    return 0.0, None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("clips_json")
    ap.add_argument("beats_json")
    ap.add_argument("--mezz-dir", default="work/mezz")
    ap.add_argument("--target-duration", type=float, default=30)
    ap.add_argument("--out", default="plan/cut_plan.json")
    ap.add_argument("--order", help="Comma-separated clip stems to force rotation order")
    ap.add_argument("--hook", help="Clip stem to open with (pick from contact sheets)")
    ap.add_argument("--close", help="Clip stem to end on (pick from contact sheets)")
    ap.add_argument("--climax", help="Clip stem to pin at the energy peak "
                                     "(the punch phase), not just hook/close")
    ap.add_argument("--music-start", type=float, default=None,
                    help="Seconds into the track where the reel starts. "
                         "Snapped to the nearest beat.")
    ap.add_argument("--music-start-lyric", metavar="TEXT",
                    help="Start the reel at the beat nearest this lyric. "
                         "Reads the .srt fetch_music.py saved beside the track.")
    ap.add_argument("--lyrics", help="Lyrics .srt (default: beside the track)")
    ap.add_argument("--lyric-occurrence", type=int,
                    help="Which match of --music-start-lyric to take, 1-based")
    ap.add_argument("--candidates", type=int, default=CANDIDATES,
                    help=f"In-points tried per segment when scoring content "
                         f"(default {CANDIDATES}). More is slower and better.")
    ap.add_argument("--workers", type=int, default=6,
                    help="Parallel frame scores (default 6)")
    ap.add_argument("--score-in-points", action="store_true",
                    help="Keep only in-points that contain a person, via Apple "
                         "Vision (macOS). Costs about 0.5s per candidate. "
                         "Needs pyobjc-framework-Vision; off by default.")
    ap.add_argument("--min-beats", type=int, default=2,
                    help="Minimum cut length in beats outside the punch phase (default 2)")
    ap.add_argument("--punch-beats", type=int, default=None,
                    help="Cut length during the punch phase (default = --min-beats; "
                         "pass 1 for sub-second flash cuts — use sparingly)")
    args = ap.parse_args()

    punch_beats = args.punch_beats or args.min_beats
    clips = [c for c in json.loads(Path(args.clips_json).read_text())["clips"]
             if "error" not in c]
    if not clips:
        sys.exit("No usable clips in clips.json")
    beats = json.loads(Path(args.beats_json).read_text())
    music_start, start_note = resolve_music_start(args, beats)
    # Anchor the grid where the reel actually starts. Prepending a bare 0.0
    # is wrong twice over: on a track with a quiet intro the first detected
    # beat can be many seconds in (12.5s on one real track), so the opening
    # hold swallows that whole gap and blows the target duration; and when
    # music_start is set, the cuts must line up with beats THERE, not with
    # beats at the top of the track.
    _start = music_start
    _after = [b for b in beats["beats"] if b >= _start - 1e-6]
    if not _after:
        sys.exit(f"No beats at or after music_start={_start:.2f}s — the track "
                 f"ends before that, or beat detection found nothing there.")
    beat_times = ([_start] + _after) if abs(_after[0] - _start) > 1e-6 else _after
    mezz = Path(args.mezz_dir)

    def mezz_path(c):
        return str(mezz / (Path(c["file"]).stem + ".mp4"))

    by_stem = {Path(c["file"]).stem: c for c in clips}
    climax_clip = by_stem.get(args.climax) if args.climax else None
    if args.climax and climax_clip is None:
        sys.exit(f"--climax {args.climax!r} is not one of the clips.")
    if args.order:
        clips = [by_stem[s] for s in args.order.split(",") if s in by_stem]
    if args.hook and args.hook in by_stem:
        clips = [by_stem[args.hook]] + [c for c in clips if c is not by_stem[args.hook]]
    close_clip = by_stem.get(args.close) if args.close else None

    spb = 60.0 / beats["bpm"]                      # seconds per beat
    hold_beats = 8 if args.target_duration >= 20 else 4   # hook/close holds
    hold_beats = min(hold_beats, max(2, (len(beat_times) - 1) // 4))

    # --- Pass 1: lay out (phase, beat-count) slots against the beat grid ---
    def beats_for(phase_mult):
        return punch_beats if phase_mult == 0 else args.min_beats * phase_mult

    body_target = max(0.0, args.target_duration - 2 * hold_beats * spb)
    slots = [("hook", hold_beats)]
    body_elapsed, i = 0.0, hold_beats
    while body_elapsed < body_target and i < len(beat_times) - 1:
        frac, cum = body_elapsed / body_target if body_target else 1.0, 0.0
        phase = BODY_PHASES[-1]
        for p in BODY_PHASES:
            cum += p[2]
            if frac < cum:
                phase = p
                break
        n = beats_for(phase[1])
        end_idx = min(i + n, len(beat_times) - 1)
        dur = beat_times[end_idx] - beat_times[i]
        if dur <= 0.05:
            break
        slots.append((phase[0], end_idx - i))
        body_elapsed += dur
        i = end_idx
    close_end = min(i + hold_beats, len(beat_times) - 1)
    if beat_times[close_end] - beat_times[i] > 0.5:
        slots.append(("close", close_end - i))

    # --- Pass 2: assign a clip to each slot (rotation, no same-clip adjacency) ---
    durations, t_idx = [], 0
    for _, n in slots:
        durations.append(round(beat_times[t_idx + n] - beat_times[t_idx], 3))
        t_idx += n

    score_content = args.score_in_points
    if score_content:
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        try:
            from detect_subjects import _vision
        except ImportError:
            _vision = lambda: False
        if not _vision():
            sys.exit("--score-in-points needs Apple Vision:\n"
                     "  pip install pyobjc-framework-Vision "
                     "pyobjc-framework-Quartz\n"
                     "macOS only. Drop the flag for positional in-points.")
    # one order per phase, so consecutive phases do not replay one sequence
    phase_names = []
    for ph, _ in slots:
        if not phase_names or phase_names[-1] != ph:
            phase_names.append(ph)
    phase_index = [phase_names.index(ph) for ph, _ in slots]
    assigned, prev, last_used = [], None, {}
    for k, (phase, _) in enumerate(slots):
        dur = durations[k]
        pick = None
        if phase == "close" and close_clip and close_clip["duration"] >= dur + 0.1 \
                and (close_clip is not prev or len(clips) == 1):
            pick = close_clip
        elif phase == "punch" and climax_clip \
                and climax_clip["duration"] >= dur + 0.1 \
                and (climax_clip is not prev or len(clips) == 1):
            pick = climax_clip
        else:
            pick = pick_clip(clips, prev, last_used, dur, phase_index[k])
        if pick is None:                            # nothing long enough: allow repeat
            pick = max((c for c in clips if c is not prev), default=clips[0],
                       key=lambda x: x["duration"])
            durations[k] = dur = min(dur, round(pick["duration"] - 0.1, 3))
        cnt = last_used.get(id(pick), (0, -1))[0]
        last_used[id(pick)] = (cnt + 1, k)
        assigned.append(pick)
        prev = pick

    # --- Pass 3: in-points, spread across the clip and scored for content ---
    usage = {}
    for k, c in enumerate(assigned):
        usage.setdefault(id(c), (c, []))[1].append(k)
    in_points = [0.0] * len(slots)
    jobs = []
    for c, ks in usage.values():
        for j, k in enumerate(ks):                 # usage j of len(ks), chronological
            cands = candidate_times(c["duration"], durations[k], j, len(ks),
                                    args.candidates if score_content else 1)
            jobs.append((k, c, cands))

    flat_segments = []
    if score_content:
        print(f"Scoring {sum(len(j[2]) for j in jobs)} candidate in-points "
              f"for content...", flush=True)

        def score_job(job):
            k, c, cands = job
            path = mezz_path(c)
            return k, cands, [content_score(path, t) for t in cands]

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            for k, cands, scores in pool.map(score_job, jobs):
                best, ok = best_candidate(scores)
                in_points[k] = cands[best]
                if not ok:
                    flat_segments.append((k, max(scores) if scores else 0.0))
    else:
        for k, c, cands in jobs:
            in_points[k] = cands[0]

    NOTES = {"hook": "HOOK — most kinetic shot here",
             "close": "CLOSE — payoff / smile / CTA"}
    segments = [{"clip": mezz_path(c), "in": in_points[k], "duration": durations[k],
                 "note": NOTES.get(slots[k][0], f"{slots[k][0]} · {slots[k][1]}-beat")}
                for k, c in enumerate(assigned)]
    total = round(sum(s["duration"] for s in segments), 2)

    plan = {
        "music": beats["audio"],
        "music_start": music_start,
        "bpm": beats["bpm"],
        "total_duration": total,
        "segments": segments,
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(plan, indent=2))

    if start_note:
        print(f"music_start: {start_note}\n")
    src = beats.get("source", "unrecorded")
    print(f"Cut plan: {len(segments)} segments, {total:.1f}s @ {beats['bpm']} BPM "
          f"({src}) -> {out}\n")
    if flat_segments:
        print(f"{len(flat_segments)}/{len(slots)} segment(s) contain no "
              f"person in any candidate. On POV or scenery footage that is "
              f"normal — the shot IS the motion. On footage that should show "
              f"someone, these are the empty cuts:", file=sys.stderr)
        for k, s in flat_segments[:8]:
            print(f"  segment {k + 1}: {Path(segments[k]['clip']).name} "
                  f"@{segments[k]['in']:.2f}s (best score {s:.2f})",
                  file=sys.stderr)
        print("", file=sys.stderr)
    if src == "fixed-grid":
        print("WARNING: built on a metronome grid, not detected beats — these "
              "cuts are NOT synced to the track. Tell the user before Gate 1.\n",
              file=sys.stderr)
    cum = 0.0
    for n, s in enumerate(segments):
        print(f"  {n:2d}  {cum:6.2f}s  {s['duration']:5.2f}s  "
              f"{Path(s['clip']).stem:<24} in={s['in']:<7} {s['note']}")
        cum += s["duration"]
    print("\nEdit plan/cut_plan.json (or ask for changes), then render a preview.")


if __name__ == "__main__":
    main()