#!/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.

Age-restricted, members-only and region-locked videos need a logged-in
session: pass --cookies-from-browser, or an exported --cookies file. Without
one yt-dlp reports a bare 403 and nothing says why.

On macOS, --cookies-from-browser chrome BLOCKS on a Keychain prompt for
Chrome Safe Storage — measured hanging indefinitely with no output when
nothing can click it. Unattended runs want an exported --cookies file
instead; safari needs no Keychain unlock. Quit Chrome first either way, since
a running Chrome holds a lock on its cookie store.

--metadata writes the uploader, title and date beside the file. That is the
actionable half of the rights reminder below: "get permission" is not useful
without knowing whose permission.

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
  fetch_music.py URL --cookies-from-browser chrome   # age-restricted
  fetch_music.py URL --video 1080 --outdir raw/      # footage, not music
"""
import argparse, json, shutil, subprocess, sys
from pathlib import Path


def auth_args(cookies, from_browser):
    """yt-dlp flags for a logged-in session, or an empty list."""
    if cookies:
        return ["--cookies", str(cookies)]
    if from_browser:
        return ["--cookies-from-browser", from_browser]
    return []


def video_format(height):
    """A yt-dlp format selector for a target height.

    Falls back down the chain rather than failing: a video with no 1080 stream
    should give its best available, not an error.
    """
    if height in (None, "best"):
        return "bestvideo*+bestaudio/best"
    return (f"bestvideo[height<={height}]+bestaudio/"
            f"best[height<={height}]/bestvideo*+bestaudio/best")


def read_credit(info_path):
    """'Uploader — Title (date)' from a yt-dlp .info.json, for attribution."""
    try:
        d = json.loads(Path(info_path).read_text())
    except (OSError, ValueError):
        return None
    who, what = d.get("uploader") or d.get("channel"), d.get("title")
    if not (who or what):
        return None
    when = d.get("upload_date") or ""
    when = f" ({when[:4]}-{when[4:6]}-{when[6:]})" if len(when) == 8 else ""
    return f"{who or 'unknown'} — {what or 'untitled'}{when}"


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")
    ap.add_argument("--video", nargs="?", const="best", metavar="HEIGHT",
                    help="Fetch VIDEO instead of audio: best, 2160, 1080, 720, "
                         "480. Rights are stricter here than for music — see "
                         "the warning this prints.")
    ap.add_argument("--cookies", metavar="FILE",
                    help="Netscape-format cookies file, for age-restricted, "
                         "members-only or region-locked videos")
    ap.add_argument("--cookies-from-browser", metavar="BROWSER",
                    help="Read cookies from a logged-in browser: chrome, "
                         "firefox, safari, edge, brave. On macOS, chrome "
                         "blocks on a Keychain prompt — use --cookies FILE "
                         "for unattended runs")
    ap.add_argument("--no-metadata", action="store_true",
                    help="Skip the .info.json — it records who to credit")
    ap.add_argument("--thumbnail", action="store_true",
                    help="Also save the poster frame (stylize_thumbnail.py "
                         "takes it from here)")
    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())

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

    # pass 1: the media itself (a failure here is fatal)
    if args.video:
        cmd = base + ["-f", video_format(args.video), "--merge-output-format", "mp4"]
    else:
        cmd = base + ["-f", "bestaudio/best", "-x", "--audio-format", args.format,
                      "--audio-quality", "0"]
    if not args.no_metadata:
        cmd += ["--write-info-json"]
    if args.thumbnail:
        cmd += ["--write-thumbnail"]
    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}).\n" + (
            "A 403 or 'Sign in to confirm' here means the video needs a "
            "logged-in session — retry with --cookies-from-browser chrome.\n"
            if not auth else
            "Cookies were supplied, so this is not an auth problem. A stale "
            "yt-dlp is the next thing to check: brew upgrade yt-dlp.\n"))

    # 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)
    want_ext = ".mp4" if args.video else f".{args.format}"
    media = [p for p in new if p.suffix == want_ext]
    subs = [p for p in new if p.suffix in (".srt", ".vtt")]
    info = [p for p in new if p.name.endswith(".info.json")]
    thumbs = [p for p in new if p.suffix in (".jpg", ".webp", ".png")]
    if not media:
        sys.exit(f"Download reported success but no {want_ext} file appeared — "
                 f"check yt-dlp output above.")

    track = media[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

    label = "Video" if args.video else "Track"
    print(f"\n{label}: {track}  ({dur:.0f}s)")
    for i in info:
        credit = read_credit(i)
        if credit:
            print(f"Source: {credit}")
        print(f"Metadata: {i}")
    for t in thumbs:
        print(f"Thumbnail: {t}  (stylize_thumbnail.py takes it from here)")
    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.")
    if args.video:
        print(f"\nNext: put it in raw/ and probe it like any clip.")
        print("Rights: footage is stricter than music. A track may be covered "
              "by a platform licence; someone else's video is not, and "
              "reposting it is a takedown or a strike. Confirm before Gate 1.")
    else:
        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()