#!/usr/bin/env python3
"""Detect beats in a music track for beat-synced cutting.

Primary: librosa beat tracking (+ downbeat estimate via onset strength).
--bpm N instead lays a FIXED METRONOME GRID. That is not beat detection: it
ignores the audio entirely, so cuts land on clock ticks rather than on the
track. Use it only when the user has explicitly accepted that trade. The
output records which was used in "source", and build_cut_plan.py surfaces it.

Output beats.json:
{
  "audio": "music/track.mp3", "bpm": 128.0, "source": "librosa", "duration": 45.2,
  "beats": [0.47, 0.94, ...],        # all beat times (s)
  "strong_beats": [0.47, 2.34, ...]  # downbeat-ish, best cut points
}

Usage:
  detect_beats.py music/track.mp3 --out work/beats.json
  detect_beats.py music/track.mp3 --bpm 120 --out work/beats.json   # no librosa
"""
import argparse, json, subprocess, sys
from pathlib import Path


def audio_duration(path):
    r = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries",
                        "format=duration", "-of", "csv=p=0", str(path)],
                       capture_output=True, text=True)
    return float(r.stdout.strip())


def fixed_grid(duration, bpm):
    step = 60.0 / bpm
    beats = []
    t = step
    while t < duration:
        beats.append(round(t, 3))
        t += step
    strong = beats[::4]  # every bar in 4/4
    return bpm, beats, strong


def librosa_beats(path):
    import librosa
    import numpy as np
    y, sr = librosa.load(str(path), mono=True)
    tempo, frames = librosa.beat.beat_track(y=y, sr=sr, units="frames")
    times = librosa.frames_to_time(frames, sr=sr)
    # downbeat heuristic: beats with highest local onset strength, ~1 per bar
    onset = librosa.onset.onset_strength(y=y, sr=sr)
    strengths = onset[np.minimum(frames, len(onset) - 1)]
    strong = []
    for i in range(0, len(times), 4):
        chunk = list(range(i, min(i + 4, len(times))))
        best = max(chunk, key=lambda j: strengths[j])
        strong.append(times[best])
    bpm = float(tempo if np.isscalar(tempo) else tempo[0])
    return round(bpm, 1), [round(float(t), 3) for t in times], [round(float(t), 3) for t in strong]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio")
    ap.add_argument("--out", default="work/beats.json")
    ap.add_argument("--bpm", type=float, help="Skip detection, use fixed BPM grid")
    args = ap.parse_args()

    duration = audio_duration(args.audio)
    if args.bpm:
        bpm, beats, strong = fixed_grid(duration, args.bpm)
        mode = "fixed-grid"
        print("WARNING: --bpm lays a metronome grid and never analyses the "
              "audio. Cuts will not land on the track's actual beats. Install "
              "librosa for real detection.", file=sys.stderr)
    else:
        try:
            bpm, beats, strong = librosa_beats(args.audio)
            mode = "librosa"
        except ImportError:
            raise SystemExit("librosa not installed. Either `pip install librosa soundfile` "
                             "or rerun with --bpm <track bpm> for a fixed grid.")

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({
        "audio": str(args.audio), "bpm": bpm, "source": mode,
        "duration": round(duration, 2),
        "beats": beats, "strong_beats": strong,
    }, indent=2))
    print(f"{mode}: {bpm} BPM, {len(beats)} beats over {duration:.1f}s -> {out}")


if __name__ == "__main__":
    main()