#!/usr/bin/env python3
"""Search YouTube, then save the transcript of every result.

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

  python search_transcripts.py "vector databases explained" --limit 10
  python search_transcripts.py "pricing" --channel @stripe --limit 20
  python search_transcripts.py "rag tutorial" --limit 25 --dry-run

Costs (successful calls only; errors are never charged):
  search across YouTube       4 credits per call (1-50 results)
  search inside one channel   1 credit per call  (1-50 results)
  transcript                  1 credit per video (no captions = 404, free)
  balance check (whoami)      free

Output goes to ./out/<query-slug>/:
  results.json         the search results, as returned
  <videoId>.txt        one transcript per video, with a title header
  index.json           what was saved, skipped or missing
Re-running skips videos that already have a .txt file.
"""

from __future__ import annotations

import argparse
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          # first fetch of a video can take 10-30 s
MAX_ATTEMPTS = 4      # for 408 / 429 / 503 / network errors
SEARCH_COST = 4
CHANNEL_SEARCH_COST = 1


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 = {}
            code = body.get("code") or body.get("error") or "error"
            if e.code in (408, 429, 503) and attempt < MAX_ATTEMPTS:
                wait = _retry_after(e.headers.get("Retry-After"), attempt)
                print(f"  {e.code}, retrying in {wait:.0f}s", file=sys.stderr)
                time.sleep(wait)
                continue
            raise ApiError(e.code, code, body.get("detail")) from None
        except (urllib.error.URLError, TimeoutError) as e:
            if attempt < MAX_ATTEMPTS:
                wait = 2 ** attempt
                print(f"  network error ({e}), retrying in {wait}s", file=sys.stderr)
                time.sleep(wait)
                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 slug(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")[:60] or "search"


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="Search YouTube and save the transcript of every result.")
    p.add_argument("query", help="what to search for")
    p.add_argument("--channel", help="search inside one channel only (@handle, URL or UC... ID): 1 credit instead of 4")
    p.add_argument("--limit", type=int, default=10, help="results to fetch, 1-50 (default 10)")
    p.add_argument("--language", default="en", help="preferred caption language (default en)")
    p.add_argument("--max-credits", type=int, default=None, help="refuse to run if the worst case costs more than this")
    p.add_argument("--out", default="out", help="output folder (default ./out)")
    p.add_argument("--dry-run", action="store_true", help="search only, show what transcripts would cost, save nothing else")
    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
    if not 1 <= a.limit <= 50:
        print("--limit must be 1-50", file=sys.stderr)
        return 2

    search_cost = CHANNEL_SEARCH_COST if a.channel else SEARCH_COST
    worst_case = search_cost + (0 if a.dry_run else a.limit)

    try:
        balance = call("/account/whoami", {}, key)["credits"]["spendable"]
    except ApiError as e:
        print(f"Key check failed: {e}", file=sys.stderr)
        return 1
    print(f"Balance: {balance} credits. Worst case for this run: {worst_case} credits.")
    if a.max_credits is not None and worst_case > a.max_credits:
        print(f"Worst case {worst_case} is over --max-credits {a.max_credits}; lower --limit.", file=sys.stderr)
        return 1
    if worst_case > balance:
        print(f"Not enough credits for the worst case ({worst_case} > {balance}); lower --limit or top up.", file=sys.stderr)
        return 1

    try:
        if a.channel:
            found = call("/youtube/channel/search", {"channel": a.channel, "q": a.query, "limit": a.limit}, key)
        else:
            found = call("/youtube/search", {"q": a.query, "type": "video", "limit": a.limit}, key)
    except ApiError as e:
        print(f"Search failed (not charged): {e}", file=sys.stderr)
        return 1

    videos = [r for r in found.get("results", []) if r.get("videoId")][: a.limit]
    print(f"Found {len(videos)} videos ({search_cost} credit{'s' if search_cost > 1 else ''}).")

    folder = Path(a.out) / slug((a.channel + "-" if a.channel else "") + a.query)
    if a.dry_run:
        for v in videos:
            print(f"  {v['videoId']}  {v.get('title', '')}")
        print(f"Transcripts would cost at most {len(videos)} credits (videos without captions are free).")
        return 0

    folder.mkdir(parents=True, exist_ok=True)
    (folder / "results.json").write_text(json.dumps(found, indent=2, ensure_ascii=False), encoding="utf-8")

    index = {"query": a.query, "channel": a.channel, "saved": [], "skipped": [], "missing": [], "credits_used": search_cost}
    for n, v in enumerate(videos, 1):
        vid, title = v["videoId"], v.get("title") or ""
        path = folder / f"{vid}.txt"
        if path.exists():
            index["skipped"].append(vid)
            print(f"[{n}/{len(videos)}] {vid} already saved, skipping")
            continue
        try:
            data = call("/youtube/transcript", {"video_url": vid, "format": "json", "include_timestamp": "false", "language": a.language}, key)
        except ApiError as e:
            if e.status == 404:
                index["missing"].append({"videoId": vid, "title": title, "reason": e.code})
                print(f"[{n}/{len(videos)}] {vid} has no transcript (free)")
                continue
            if e.status == 402:
                print(f"Out of credits after {len(index['saved'])} transcripts; stopping. {e.detail}", file=sys.stderr)
                break
            index["missing"].append({"videoId": vid, "title": title, "reason": f"{e.status} {e.code}"})
            print(f"[{n}/{len(videos)}] {vid} failed: {e}", file=sys.stderr)
            continue
        header = f"{title}\nhttps://www.youtube.com/watch?v={vid}\nchannel: {v.get('channelTitle', '')}\nlanguage: {data.get('language', '')}\n\n"
        path.write_text(header + transcript_text(data) + "\n", encoding="utf-8")
        index["saved"].append({"videoId": vid, "title": title, "file": path.name})
        index["credits_used"] += 1
        print(f"[{n}/{len(videos)}] {vid} saved")

    (folder / "index.json").write_text(json.dumps(index, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"Done: {len(index['saved'])} saved, {len(index['skipped'])} skipped, {len(index['missing'])} without transcript. "
          f"About {index['credits_used']} credits used. Files in {folder}")
    return 0


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