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

Primary: librosa beat tracking (+ downbeat estimate via onset strength).
Fallback: --bpm N generates a fixed grid (no librosa needed).

Output beats.json:
{
  "audio": "music/track.mp3", "bpm": 128.0, "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
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"
    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, "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()