#!/usr/bin/env python3
"""Fetch a music track (and subtitles/lyrics, if available) from YouTube.

The fallback for Step 4 when the user has no music file yet: they paste a
YouTube link, this pulls the audio into music/ ready for detect_beats.py.
Subtitles (manual or auto-generated) are saved as .srt next to the track —
useful for picking music_start at a specific lyric, or for caption burn-in
via render_reel.py --srt.

Requires yt-dlp (brew install yt-dlp). Downloads are atomic (yt-dlp writes
.part files and renames on completion). Filenames are sanitized to be
ffmpeg-path-safe.

Rights reminder: only use tracks the user has rights to publish, or that the
target platform licenses (TikTok/IG in-app libraries don't cover externally
rendered audio). Surface this to the user — a muted or taken-down reel wastes
the whole pipeline.

Usage:
  fetch_music.py "https://youtube.com/watch?v=..." --outdir music/
  fetch_music.py URL --section 1:10-1:50      # just the drop, not all 4 min
  fetch_music.py URL --no-subs --format mp3
"""
import argparse, shutil, subprocess, sys
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("url", help="YouTube (or any yt-dlp-supported) URL")
    ap.add_argument("--outdir", default="music")
    ap.add_argument("--format", default="m4a", choices=["m4a", "mp3", "wav"],
                    help="Audio container (default m4a — smallest, ffmpeg-native)")
    ap.add_argument("--section", metavar="START-END",
                    help="Download only this time range, e.g. 1:10-1:50 "
                         "(full songs are 3-4 min; a reel needs less)")
    ap.add_argument("--sub-langs", default="en,en-orig",
                    help="Subtitle languages (default 'en,en-orig'; exact codes — "
                         "wildcards like en.* pull dozens of auto-translations "
                         "and get rate-limited by YouTube)")
    ap.add_argument("--no-subs", action="store_true",
                    help="Skip subtitles/lyrics")
    args = ap.parse_args()

    if not shutil.which("yt-dlp"):
        sys.exit("yt-dlp not found — install with: brew install yt-dlp")

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    before = set(outdir.iterdir())

    base = ["yt-dlp", "--no-playlist", "--restrict-filenames",
            "-o", str(outdir / "%(title)s.%(ext)s")]

    # pass 1: audio (the deliverable — a failure here is fatal)
    cmd = base + ["-f", "bestaudio/best", "-x", "--audio-format", args.format,
                  "--audio-quality", "0"]
    if args.section:
        # force-keyframes makes the cut exact; without it yt-dlp trims on
        # stream fragment boundaries (several seconds of slop)
        cmd += ["--download-sections", f"*{args.section}",
                "--force-keyframes-at-cuts"]
    r = subprocess.run(cmd + [args.url])
    if r.returncode != 0:
        sys.exit(f"yt-dlp failed (exit {r.returncode})")

    # pass 2: subtitles (best-effort — never fail the run over lyrics)
    if not args.no_subs:
        r = subprocess.run(base + ["--skip-download", "--write-subs",
                                   "--write-auto-subs", "--sub-langs",
                                   args.sub_langs, "--convert-subs", "srt",
                                   args.url])
        if r.returncode != 0:
            print("WARNING: subtitle fetch failed — continuing without lyrics",
                  file=sys.stderr)

    new = sorted(set(outdir.iterdir()) - before)
    audio = [p for p in new if p.suffix == f".{args.format}"]
    subs = [p for p in new if p.suffix in (".srt", ".vtt")]
    if not audio:
        sys.exit("Download reported success but no audio file appeared — "
                 "check yt-dlp output above.")

    track = audio[0]
    p = subprocess.run(["ffprobe", "-v", "quiet", "-show_entries",
                        "format=duration", "-of", "csv=p=0", str(track)],
                       capture_output=True, text=True)
    dur = float(p.stdout.strip()) if p.stdout.strip() else 0.0

    print(f"\nTrack: {track}  ({dur:.0f}s)")
    for s in subs:
        print(f"Subtitles: {s}  (lyrics for music_start picks, or --srt burn-in)")
    if not subs and not args.no_subs:
        print("No subtitles available for this video.")
    print(f"\nNext: python3 scripts/detect_beats.py '{track}' --out work/beats.json")
    print("Reminder: confirm the user has rights to use this track on the "
          "target platforms.")


if __name__ == "__main__":
    main()