#!/bin/sh
"exec" "`dirname $0`/../venvs/index/bin/python" "$0" "$@"
"""vg-beats — extract a beat grid from a music track for cutting in Resolve.

Usage: vg-beats SONG.mp3 [--fps 25] [--tc-start 01:00:00:00] [--bpm HINT] [--out PREFIX]

Writes:
  PREFIX.beats.csv   one beat time (seconds) per line, with beat number
  PREFIX.beats.edl   Resolve timeline markers — import via
                     Timeline > (right-click) > Timelines > Import > Timeline Markers from EDL
Every 4th beat is a red marker (bar guess), others blue.
"""
import argparse, os, sys

os.environ["PATH"] = "/opt/homebrew/bin:" + os.environ.get("PATH", "")


def tc(seconds, fps, start_frames):
    total = int(round(seconds * fps)) + start_frames
    f = total % int(fps)
    s = (total // int(fps)) % 60
    m = (total // (int(fps) * 60)) % 60
    h = total // (int(fps) * 3600)
    return f"{h:02d}:{m:02d}:{s:02d}:{f:02d}"


def parse_tc(t, fps):
    h, m, s, f = (int(x) for x in t.split(":"))
    return ((h * 3600 + m * 60 + s) * int(fps)) + f


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("audio")
    ap.add_argument("--fps", type=float, default=25.0, help="timeline fps for the EDL (default 25)")
    ap.add_argument("--tc-start", default="01:00:00:00", help="timeline start timecode (Resolve default 01:00:00:00)")
    ap.add_argument("--bpm", type=float, default=None, help="starting tempo hint")
    ap.add_argument("--out", default=None, help="output prefix (default: alongside the audio file)")
    args = ap.parse_args()

    import librosa
    import numpy as np

    print(f"loading {args.audio} ...")
    y, sr = librosa.load(args.audio, sr=22050, mono=True)
    kw = {"start_bpm": args.bpm} if args.bpm else {}
    tempo, beats = librosa.beat.beat_track(y=y, sr=sr, units="time", **kw)
    tempo = float(np.atleast_1d(tempo)[0])
    if len(beats) == 0:
        sys.exit("no beats detected — is this a music track?")
    print(f"tempo ~{tempo:.1f} BPM, {len(beats)} beats over {len(y)/sr:.1f}s")

    prefix = args.out or os.path.splitext(args.audio)[0]
    start_frames = parse_tc(args.tc_start, args.fps)

    with open(prefix + ".beats.csv", "w") as f:
        f.write("beat,seconds\n")
        for i, t in enumerate(beats, 1):
            f.write(f"{i},{t:.4f}\n")

    with open(prefix + ".beats.edl", "w") as f:
        f.write(f"TITLE: {os.path.basename(prefix)} beats ({tempo:.1f}bpm)\nFCM: NON-DROP FRAME\n\n")
        for i, t in enumerate(beats, 1):
            a = tc(t, args.fps, start_frames)
            b = tc(t + 1.0 / args.fps, args.fps, start_frames)
            color = "ResolveColorRed" if (i - 1) % 4 == 0 else "ResolveColorBlue"
            f.write(f"{i:03d}  001      V     C        {a} {b} {a} {b}  \n")
            f.write(f" |C:{color} |M:Beat {i} |D:1\n\n")

    print(f"wrote {prefix}.beats.csv and {prefix}.beats.edl")
    print("Resolve: right-click the timeline in the media pool > Timelines > Import > Timeline Markers from EDL")


if __name__ == "__main__":
    main()