#!/usr/bin/env python3
"""Captions first: get the transcript of every video that has captions, and list
the ones that would need speech-to-text, with an estimate of what that would cost.

Uses YouTubeTranscriptAPI (https://youtubetranscriptapi.com). Standard library
only, Python 3.9+. Set YT_API_KEY first.

  python captions_first.py dQw4w9WgXcQ https://youtu.be/abc123xyz00
  python captions_first.py --file video_ids.txt
  python captions_first.py --playlist PLxxxxxxxx
  python captions_first.py --channel @TED --stt-rate 0.0043

Costs (successful calls only; errors are never charged):
  transcript                    1 credit per video (no captions = 404, free)
  playlist or channel listing   20 credits (up to 100 videos)
  balance check (whoami)        free

Output in ./captions_first/:
  <videoId>.txt          transcript text, with a title header when known
  needs_stt.csv          videos with no captions: id, title, length, est. minutes, est. cost
  summary.json           counts, credits used, speech-to-text estimate
Re-running skips videos that already have a .txt file.

This script never downloads audio. For videos without captions it only tells you
what speech-to-text would cost; check that you have the rights to the audio first.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

BASE = os.environ.get("YT_API_BASE", "https://youtubetranscriptapi.com/api/v2")
TIMEOUT = 60
MAX_ATTEMPTS = 4
LISTING_COST = 20
DEFAULT_STT_RATE = 0.0045  # USD per audio minute; OpenAI gpt-transcribe list price, Nov 2026
ID_RE = re.compile(r"(?:v=|youtu\.be/|shorts/|embed/|live/)([A-Za-z0-9_-]{11})|^([A-Za-z0-9_-]{11})$")


class ApiError(Exception):
    def __init__(self, status: int, code: str, detail):
        super().__init__(f"{status} {code}: {detail}")
        self.status, self.code, self.detail = status, code, detail


def call(path: str, params: dict, key: str) -> dict:
    """GET one endpoint. Retries 408/429/503 and network errors, honoring Retry-After."""
    url = f"{BASE}{path}?{urllib.parse.urlencode(params)}"
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {key}", "Accept": "application/json"})
    for attempt in range(1, MAX_ATTEMPTS + 1):
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            try:
                body = json.loads(e.read().decode("utf-8") or "{}")
            except ValueError:
                body = {}
            if e.code in (408, 429, 503) and attempt < MAX_ATTEMPTS:
                time.sleep(_retry_after(e.headers.get("Retry-After"), attempt))
                continue
            raise ApiError(e.code, body.get("code") or body.get("error") or "error", body.get("detail")) from None
        except (urllib.error.URLError, TimeoutError):
            if attempt < MAX_ATTEMPTS:
                time.sleep(2 ** attempt)
                continue
            raise
    raise RuntimeError("unreachable")


def _retry_after(value: str | None, attempt: int) -> float:
    try:
        return max(1.0, min(float(value), 120.0))
    except (TypeError, ValueError):
        return float(2 ** attempt)


def video_id(text: str) -> str | None:
    m = ID_RE.search(text.strip())
    return (m.group(1) or m.group(2)) if m else None


def minutes(length_text: str | None) -> float | None:
    """'1:02:03' -> 62.05, '3:33' -> 3.55; None when unknown."""
    if not length_text or not re.fullmatch(r"\d+(:\d{1,2}){1,2}", length_text.strip()):
        return None
    secs = 0
    for part in length_text.strip().split(":"):
        secs = secs * 60 + int(part)
    return round(secs / 60, 2)


def transcript_text(data: dict) -> str:
    t = data.get("transcript")
    if isinstance(t, str):
        return t
    return "\n".join(seg.get("text", "") for seg in t or [])


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Captions first; list the videos that would need speech-to-text.")
    p.add_argument("videos", nargs="*", help="video IDs or URLs")
    p.add_argument("--file", help="text file with one video ID or URL per line")
    p.add_argument("--playlist", help="playlist ID or URL (20 credits, up to 100 videos)")
    p.add_argument("--channel", help="channel @handle, URL or UC ID: its 100 newest uploads (20 credits)")
    p.add_argument("--language", default="en", help="preferred caption language (default en)")
    p.add_argument("--stt-rate", type=float, default=DEFAULT_STT_RATE, help=f"speech-to-text USD per minute for the estimate (default {DEFAULT_STT_RATE})")
    p.add_argument("--out", default="captions_first", help="output folder")
    a = p.parse_args(argv)

    key = os.environ.get("YT_API_KEY")
    if not key:
        print("Set YT_API_KEY first (free key, 100 credits: https://youtubetranscriptapi.com/signup)", file=sys.stderr)
        return 2

    # Collect videos: (id, title, lengthText)
    videos: list[tuple[str, str, str | None]] = []
    raw = list(a.videos)
    if a.file:
        raw += [ln for ln in Path(a.file).read_text(encoding="utf-8").splitlines() if ln.strip() and not ln.startswith("#")]
    for item in raw:
        vid = video_id(item)
        if vid:
            videos.append((vid, "", None))
        else:
            print(f"Skipping, not a video ID or URL: {item}", file=sys.stderr)

    credits = 0
    try:
        balance = call("/account/whoami", {}, key)["credits"]["spendable"]
        for kind, value in (("playlist", a.playlist), ("channel", a.channel)):
            if not value:
                continue
            listing = call("/youtube/playlist/videos" if kind == "playlist" else "/youtube/channel/videos", {kind: value}, key)
            credits += LISTING_COST
            for r in listing.get("results", []):
                if r.get("videoId"):
                    videos.append((r["videoId"], r.get("title") or "", r.get("lengthText")))
            if listing.get("has_more"):
                print(f"Note: the {kind} has more than the {len(listing.get('results', []))} videos listed here.")
    except ApiError as e:
        print(f"Request failed (not charged): {e}", file=sys.stderr)
        return 1

    seen, unique = set(), []
    for v in videos:
        if v[0] not in seen:
            seen.add(v[0])
            unique.append(v)
    if not unique:
        print("No videos to process.", file=sys.stderr)
        return 2
    print(f"{len(unique)} videos. Balance {balance} credits; transcripts cost at most {len(unique)}.")

    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    saved, skipped, needs, failed = [], [], [], []
    for n, (vid, title, length) in enumerate(unique, 1):
        path = out / f"{vid}.txt"
        if path.exists():
            skipped.append(vid)
            continue
        try:
            data = call("/youtube/transcript", {"video_url": vid, "include_timestamp": "false",
                                                "send_metadata": "true", "language": a.language}, key)
        except ApiError as e:
            if e.status == 404:
                needs.append((vid, title, length))
                print(f"[{n}/{len(unique)}] {vid}: no captions (free)")
                continue
            if e.status == 402:
                print(f"Out of credits after {len(saved)} transcripts; stopping.", file=sys.stderr)
                break
            failed.append({"videoId": vid, "error": f"{e.status} {e.code}"})
            print(f"[{n}/{len(unique)}] {vid}: {e}", file=sys.stderr)
            continue
        credits += 1
        meta = data.get("metadata") or {}
        title = title or meta.get("title") or ""
        header = f"{title}\nhttps://www.youtube.com/watch?v={vid}\nlanguage: {data.get('language', '')}\n\n"
        path.write_text(header + transcript_text(data) + "\n", encoding="utf-8")
        saved.append(vid)
        print(f"[{n}/{len(unique)}] {vid}: saved")

    est_minutes = 0.0
    unknown = 0
    with open(out / "needs_stt.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["videoId", "title", "length", "est_minutes", "est_cost_usd"])
        for vid, title, length in needs:
            m = minutes(length)
            if m is None:
                unknown += 1
            else:
                est_minutes += m
            w.writerow([vid, title, length or "", "" if m is None else m, "" if m is None else round(m * a.stt_rate, 4)])

    summary = {
        "videos": len(unique), "saved": len(saved), "skipped_existing": len(skipped),
        "no_captions": len(needs), "failed": failed, "credits_used": credits,
        "stt_estimate": {"rate_usd_per_min": a.stt_rate, "known_minutes": round(est_minutes, 1),
                         "usd": round(est_minutes * a.stt_rate, 2), "videos_with_unknown_length": unknown},
    }
    (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"\nSaved {len(saved)}, skipped {len(skipped)}, no captions {len(needs)}, failed {len(failed)}. "
          f"Credits used: {credits}.")
    if needs:
        extra = f" (+{unknown} of unknown length)" if unknown else ""
        print(f"Speech-to-text for the rest: ~{est_minutes:.0f} min{extra}, about ${est_minutes * a.stt_rate:.2f} "
              f"at ${a.stt_rate}/min. List: {out / 'needs_stt.csv'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())