#!/usr/bin/env python3
"""Download a route video, sample frames every N seconds, and build labelled
contact sheets. Built for Daikoku entrance/exit reference gathering.

Usage: route.py <videoid> <shortname> [step_seconds]
"""
import os, re, subprocess, sys, glob, json
from PIL import Image, ImageDraw

YTDLP = os.path.expanduser("~/.local/bin/yt-dlp")
ROOT = "/tmp/routes"


def run(cmd, **kw):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True, **kw)


def meta(vid):
    r = run(f'{YTDLP} --no-warnings -J --skip-download "https://www.youtube.com/watch?v={vid}" 2>/dev/null')
    try:
        d = json.loads(r.stdout)
    except Exception:
        return {}
    return {k: d.get(k) for k in ("title", "uploader", "duration", "upload_date", "width", "height")}


def fetch(vid, name):
    d = f"{ROOT}/{name}"
    os.makedirs(d, exist_ok=True)
    mp4 = f"{d}/{name}.mp4"
    if not os.path.exists(mp4):
        r = run(f'{YTDLP} -f "bv*[height<=1080][ext=mp4]/bv*[height<=1080]/b" '
                f'--no-warnings -o "{mp4}" "https://www.youtube.com/watch?v={vid}"')
        if not os.path.exists(mp4):
            print(f"  DOWNLOAD FAILED: {r.stderr[-300:]}")
            return None
    return mp4


def sample(mp4, d, step):
    dur = float(run(f'ffprobe -v error -show_entries format=duration '
                    f'-of default=nw=1:nk=1 "{mp4}"').stdout.strip() or 0)
    for f in glob.glob(f"{d}/fr_*.jpg"):
        os.remove(f)
    ts = list(range(2, int(dur), step))
    for t in ts:
        run(f'ffmpeg -v error -ss {t} -i "{mp4}" -frames:v 1 -q:v 2 '
            f'-vf scale=960:-1 "{d}/fr_{t:04d}.jpg" -y')
    return dur, ts


def sheets(d, name, cols=3, per=9):
    fs = sorted(glob.glob(f"{d}/fr_*.jpg"))
    out = []
    for page in range((len(fs) + per - 1) // per):
        sub = fs[page * per:(page + 1) * per]
        ims = [Image.open(f).convert("RGB") for f in sub]
        w, h = ims[0].size
        rows = (len(sub) + cols - 1) // cols
        s = Image.new("RGB", (w * cols, h * rows), (10, 10, 10))
        for i, im in enumerate(ims):
            s.paste(im, ((i % cols) * w, (i // cols) * h))
        dr = ImageDraw.Draw(s)
        for i, f in enumerate(sub):
            t = int(re.search(r"fr_(\d+)", f).group(1))
            x, y = (i % cols) * w, (i // cols) * h
            dr.rectangle([x + 6, y + 6, x + 190, y + 54], fill=(0, 0, 0))
            dr.text((x + 20, y + 22), f"t={t}s", fill=(255, 225, 0))
        s = s.resize((int(s.width * 0.60), int(s.height * 0.60)), Image.LANCZOS)
        o = f"{d}/{name}_sheet{page+1}.jpg"
        s.save(o, quality=90)
        out.append(o)
    return out


if __name__ == "__main__":
    vid, name = sys.argv[1], sys.argv[2]
    step = int(sys.argv[3]) if len(sys.argv) > 3 else 8
    m = meta(vid)
    print(f"=== {name} [{vid}] ===")
    print(f"  title: {m.get('title')}")
    print(f"  dur  : {m.get('duration')}s  {m.get('width')}x{m.get('height')}")
    mp4 = fetch(vid, name)
    if not mp4:
        sys.exit(1)
    d = f"{ROOT}/{name}"
    dur, ts = sample(mp4, d, step)
    sh = sheets(d, name)
    print(f"  frames: {len(ts)} @ {step}s  sheets: {len(sh)}")
    for s in sh:
        print(f"   {s}")