#!/bin/zsh
# baton-yt — YouTube ingestion for the video-comprehension pipeline.
#
# Grabs content with yt-dlp and feeds the same decomposition every other
# footage source gets (see the video-comprehension skill): audio -> Apple
# on-device STT transcript, video -> ffmpeg contact sheet, on-screen text ->
# Apple OCR timeline. So any YouTube URL becomes readable text + one
# glanceable image — no video "watched", no cloud model needed.
#
# PATH-dispatched baton module: `baton yt ...` finds this via the git-style
# module contract; --describe feeds `baton help`.
#
#   baton yt audio <url>              audio only (m4a) -> cache dir
#   baton yt video <url>              720p mp4 -> cache dir
#   baton yt brief <url> [--sheets]   transcript: creator subtitle track if one
#                                     exists (ground truth), else Apple STT
#                                     (+ video -> contact sheet with --sheets)
#   baton yt ocr <url>                on-screen text timeline: frames every 2s
#                                     -> Apple OCR -> timecoded unique blocks
#                                     (the citation layer narration never reads)
#   baton yt subs <url> [--lang xx]   YouTube's own subtitle track as timecoded
#                                     text (creator track preferred, auto-caption
#                                     fallback) — no media download, no STT
#   Cache: ~/.baton/var/yt/<title-slug>--<video-id>/ (re-runs reuse downloads + OCR)
set -e -u -o pipefail

describe="yt — YouTube ingestion: audio/video grab + transcript/contact-sheet/OCR-timeline/subs brief (video-comprehension pipeline)"
[[ "${1:-}" == "--describe" ]] && { print -r -- "$describe"; exit 0; }

verb="${1:-}"; shift 2>/dev/null || true
url="${1:-}"; shift 2>/dev/null || true
[[ -z "$verb" || "$verb" == "--help" || "$verb" == "-h" ]] && {
  sed -n '2,24p' "$0" | sed 's/^# \{0,1\}//'; exit 0; }
[[ -z "$url" ]] && { echo "baton-yt: $verb needs a URL" >&2; exit 2; }

# Title -> filesystem-safe slug. The video id still carries identity, so this
# only has to be readable: transliterate to ASCII, fold to lowercase, collapse
# everything else to single dashes, and cap the length so the slug plus the id
# stays well inside the 255-byte filename limit.
slugify() {
  local s
  s=$(print -r -- "$1" | iconv -c -t 'ascii//TRANSLIT' 2>/dev/null || true)
  [[ -n "$s" ]] || s="$1"
  s="${s:l}"
  s=$(print -r -- "$s" | sed -E 's/[^a-z0-9]+/-/g; s/^-+//; s/-+$//')
  s="${s[1,60]}"
  print -r -- "${s%-}"
}

# One metadata call serves both the folder name and the brief header. Separate
# --print lines rather than one delimited line: titles routinely contain the
# separator ("WWDC23: … | Apple") but never a newline. On a playlist URL yt-dlp
# repeats the block per entry, so the first four lines are the first video.
meta_lines=("${(@f)$(yt-dlp --no-warnings --print '%(id)s' --print '%(title)s' \
  --print '%(uploader)s' --print '%(duration_string)s' "$url" 2>/dev/null)}")
vid="${meta_lines[1]:-}"
title="${meta_lines[2]:-}"
uploader="${meta_lines[3]:-}"
duration="${meta_lines[4]:-}"
[[ -n "$vid" ]] || { echo "baton-yt: cannot resolve $url" >&2; exit 1; }

# A folder is named <title-slug>--<id>, but the ID alone is the identity: the
# slug is cosmetic and a title can be edited after upload. So resolve an
# existing cache by id first (renaming it up to the current title), and only
# fall back to a bare id when there is no title to slug.
root="$HOME/.baton/var/yt"
slug=$(slugify "$title")
existing=("$root"/*"--$vid"(N))
if [[ -n "$slug" ]]; then
  dir="$root/$slug--$vid"
  if [[ ! -d "$dir" ]]; then
    if (( $#existing )); then
      mv "${existing[1]}" "$dir"          # title changed since the last run
    elif [[ -d "$root/$vid" ]]; then
      mv "$root/$vid" "$dir"              # folder predates title naming
    fi
  fi
elif (( $#existing )); then
  dir="${existing[1]}"                    # keep the readable name when the lookup fails
else
  dir="$root/$vid"
fi
mkdir -p "$dir"

fetch_audio() {
  local out="$dir/audio.m4a"
  [[ -s "$out" ]] || yt-dlp -q -f "bestaudio[ext=m4a]/bestaudio" -o "$out" "$url"
  print -r -- "$out"
}

fetch_video() {
  local out="$dir/video.mp4"
  [[ -s "$out" ]] || yt-dlp -q -f "best[height<=720][ext=mp4]/best[ext=mp4]/best" -o "$out" "$url"
  print -r -- "$out"
}

# Ensure the <lang> subtitle track json3 is cached; prints "<kind> <path>"
# (kind: creator|auto) or returns 1 when the video has no such track. Creator
# tracks are ground truth; auto-captions are STT-grade (name misreads, [ __ ]
# profanity censoring), so provenance is kept in the cached filename.
fetch_subs() {
  local lang="$1"
  local jc="$dir/subs.$lang.creator.json3" ja="$dir/subs.$lang.auto.json3"
  [[ -s "$jc" ]] && { print -r -- "creator $jc"; return 0; }
  [[ -s "$ja" ]] && { print -r -- "auto $ja"; return 0; }
  local raw="$dir/subs.$lang.json3"
  yt-dlp -q --no-warnings --skip-download --write-subs \
    --sub-langs "$lang" --sub-format json3 -o "$dir/subs" "$url" || true
  [[ -s "$raw" ]] && { mv "$raw" "$jc"; print -r -- "creator $jc"; return 0; }
  yt-dlp -q --no-warnings --skip-download --write-auto-subs \
    --sub-langs "$lang" --sub-format json3 -o "$dir/subs" "$url" || true
  [[ -s "$raw" ]] && { mv "$raw" "$ja"; print -r -- "auto $ja"; return 0; }
  return 1
}

# json3 caption events -> "[MM:SS] line" cues on stdout.
subs_to_cues() {
  python3 - "$1" <<'EOF'
import sys, json

def ts(ms):
    s = ms // 1000
    return f"{s//60:02d}:{s%60:02d}"

for e in json.load(open(sys.argv[1]))['events']:
    text = ''.join(seg.get('utf8', '') for seg in e.get('segs', [])).strip()
    if text:   # window-setup and aAppend newline events carry no words
        print(f"[{ts(e['tStartMs'])}] {' '.join(text.split())}")
EOF
}

case "$verb" in
  audio) fetch_audio ;;
  video) fetch_video ;;
  brief)
    sheets=0
    [[ "${1:-}" == "--sheets" ]] && sheets=1
    # Title + uploader for the header, from the metadata already resolved above.
    meta=""
    [[ -n "$title" ]] && meta="$title | $uploader | $duration"
    t="$dir/transcript.txt"
    if [[ ! -s "$t" ]]; then
      # A creator-uploaded track beats STT (it's what was actually said, with
      # timecodes, no audio download). Auto-captions do NOT displace STT —
      # they carry the same misread class plus censoring; 'subs' serves those.
      if info=$(fetch_subs en) && [[ "${info%% *}" == "creator" ]]; then
        { [[ -n "$meta" ]] && print -r -- "# $meta"
          print -r -- "# source: creator subtitle track"
          subs_to_cues "${info#* }"; } > "$t"
      else
        a=$(fetch_audio)
        wav="$dir/audio-16k.wav"
        [[ -s "$wav" ]] || ffmpeg -v error -y -i "$a" -ac 1 -ar 16000 "$wav"
        { [[ -n "$meta" ]] && print -r -- "# $meta"
          print -r -- "# source: apple stt"
          baton apple stt "$wav"; } > "$t"
      fi
    fi
    if (( sheets )); then
      v=$(fetch_video)
      s="$dir/sheet.jpg"
      if [[ ! -s "$s" ]]; then
        d=$(ffprobe -v quiet -show_entries format=duration -of csv=p=0 "$v")
        ffmpeg -v error -y -i "$v" -vf "fps=30/$d,scale=320:-1,tile=6x5" -frames:v 1 "$s"
      fi
      print -r -- "sheet: $s"
    fi
    print -r -- "transcript: $t"
    ;;
  ocr)
    # Videos cite sources visually — X-post screenshots, slides, headlines —
    # that the narration never reads aloud, so the transcript alone misses a
    # whole citation layer. This extracts it: sample frames, OCR each on-device,
    # merge near-identical reads into timecoded spans.
    v=$(fetch_video)
    tl="$dir/ocr-timeline.md"
    if [[ ! -s "$tl" ]]; then
      frames="$dir/frames"
      mkdir -p "$frames"
      # 1 frame every 2 seconds; deterministic coverage beats scene detection
      # (slow overlays fade in without a cut, so scene filters miss them)
      [[ -s "$frames/00001.png" ]] || ffmpeg -v error -i "$v" -vf fps=1/2 "$frames/%05d.png"
      # Per-frame OCR is resumable: empty .txt = prior failure, retried here.
      for f in "$frames"/*.png; do
        txt="${f%.png}.txt"
        [[ -s "$txt" ]] || baton apple ocr "$f" > "$txt" 2>/dev/null || : > "$txt"
      done
      python3 - "$frames" "$dir" <<'EOF'
import sys, os, re, difflib
frames_dir, out_dir = sys.argv[1], sys.argv[2]

def norm(t):
    return re.sub(r'\s+', ' ', t.lower()).strip()

entries = []  # (seconds, raw, normed)
for name in sorted(os.listdir(frames_dir)):
    if not name.endswith('.txt'): continue
    idx = int(name[:-4])
    secs = (idx - 1) * 2   # fps=1/2, frame 1 ~ t=0
    raw = open(os.path.join(frames_dir, name)).read().strip()
    n = norm(raw)
    if len(n) < 20:        # talking head / watermark noise
        continue
    entries.append((secs, raw, n))

# Merge consecutive near-identical texts into spans. Known limit: heavy OCR
# noise on identical screens can defeat the 0.85 gate — treat adjacent
# same-looking blocks as one when reading the output.
spans = []
for secs, raw, n in entries:
    if spans:
        prev = spans[-1]
        sim = difflib.SequenceMatcher(None, prev['norm'], n).ratio()
        if sim > 0.85:
            prev['end'] = secs
            if len(n) > len(prev['norm']):   # keep the fullest OCR read
                prev['norm'], prev['raw'] = n, raw
            continue
    spans.append({'start': secs, 'end': secs, 'raw': raw, 'norm': n})

# A span near-identical to an EARLIER one is a re-shown quote: track its
# timecodes on the original instead of duplicating the text.
uniq = []
for s in spans:
    dup_of = None
    for u in uniq:
        if difflib.SequenceMatcher(None, u['norm'], s['norm']).ratio() > 0.85:
            dup_of = u; break
    if dup_of:
        dup_of.setdefault('also', []).append((s['start'], s['end']))
    else:
        uniq.append(s)

def ts(sec):
    return f"{sec//60:02d}:{sec%60:02d}"

with open(os.path.join(out_dir, 'ocr-timeline.md'), 'w') as f:
    f.write(f"# On-screen text timeline\n\n{len(uniq)} unique text blocks from {len(entries)} text-bearing frames\n\n")
    for s in uniq:
        span = f"[{ts(s['start'])}–{ts(s['end'])}]" if s['end'] != s['start'] else f"[{ts(s['start'])}]"
        f.write(f"## {span}\n\n```\n{s['raw']}\n```\n")
        if 'also' in s:
            f.write("re-shown: " + ", ".join(f"{ts(a)}–{ts(b)}" for a,b in s['also']) + "\n")
        f.write("\n")
print(f"unique blocks: {len(uniq)} (from {len(entries)} text frames)", file=sys.stderr)
EOF
    fi
    print -r -- "ocr-timeline: $tl"
    ;;
  subs)
    lang="en"
    [[ "${1:-}" == "--lang" && -n "${2:-}" ]] && lang="$2"
    out="$dir/subs.$lang.txt"
    if [[ ! -s "$out" ]]; then
      info=$(fetch_subs "$lang") || { echo "baton-yt: no '$lang' subtitle track (creator or auto) for $url" >&2; exit 1; }
      kind="${info%% *}"
      cues=$(subs_to_cues "${info#* }")
      n=$(print -r -- "$cues" | wc -l | tr -d ' ')
      { print -r -- "# subtitles ($kind, $lang) — $n cues"
        print -r -- "$cues"; } > "$out"
      echo "$kind track, $n cues" >&2
    fi
    print -r -- "subs: $out"
    ;;
  *) echo "baton-yt: unknown verb '$verb' (audio|video|brief|ocr|subs)" >&2; exit 2 ;;
esac
