"""Download transcripts for a YouTube channel's videos with YouTubeTranscriptAPI.

    pip install httpx
    export YTT_API_KEY=tk_live_...
    python channel_transcripts.py @TED --dry-run
    python channel_transcripts.py @TED --search "climate" --search "education" --csv ted.csv

What it does:
  1. Resolves the @handle, URL or UC... ID to a channel ID (free).
  2. Lists the channel's most recent uploads, up to 100 (20 credits).
  3. Optionally runs topic searches inside the channel to reach older videos
     (1 credit each, up to 50 results per search).
  4. Fetches each transcript (1 credit per successful transcript) with a small
     concurrency limit, and appends one JSON line per video to OUT.jsonl.
     Re-running skips videos already in the file, so an interrupted run resumes.

Failed calls, 404s (no captions) and 429s are never charged.
"""

import argparse
import asyncio
import csv
import json
import os
import sys
from pathlib import Path

import httpx

API = os.environ.get("YTT_API_BASE", "https://youtubetranscriptapi.com/api/v2")
RETRYABLE = {408, 429, 500, 503}


class OutOfCredits(Exception):
    pass


async def call(client: httpx.AsyncClient, path: str, params: dict, tries: int = 5) -> httpx.Response:
    """GET with retries on transient errors, honouring Retry-After."""
    for attempt in range(tries):
        resp = await client.get(f"{API}{path}", params=params)
        if resp.status_code == 402:
            raise OutOfCredits(resp.json().get("detail", "out of credits"))
        if resp.status_code not in RETRYABLE:
            return resp
        wait = float(resp.headers.get("Retry-After", 2 ** attempt))
        await asyncio.sleep(min(wait, 60))
    return resp


async def list_videos(client: httpx.AsyncClient, channel: str, searches: list[str]) -> list[dict]:
    resolved = await call(client, "/youtube/channel/resolve", {"input": channel})
    resolved.raise_for_status()
    channel_id = resolved.json()["channel_id"]

    listing = await call(client, "/youtube/channel/videos", {"channel": channel_id})
    listing.raise_for_status()
    body = listing.json()
    videos = {v["videoId"]: v.get("title") for v in body["results"]}
    if body.get("has_more"):
        print(f"Channel has more than {len(videos)} uploads; the listing covers the newest {len(videos)}.")

    for q in searches:
        found = await call(client, "/youtube/channel/search", {"channel": channel_id, "q": q, "limit": 50})
        found.raise_for_status()
        for v in found.json()["results"]:
            if v.get("hasCaptions") is False:
                continue
            videos.setdefault(v["videoId"], v.get("title"))

    return [{"video_id": vid, "title": title} for vid, title in videos.items()]


async def fetch_one(client, sem, video, out, language, counts):
    async with sem:
        resp = await call(client, "/youtube/transcript", {"video_url": video["video_id"], "language": language})
    if resp.status_code == 404:
        counts["no_captions"] += 1
        return
    if resp.status_code != 200:
        print(f"{video['video_id']}: HTTP {resp.status_code}, skipped (re-run to retry)", file=sys.stderr)
        counts["failed"] += 1
        return
    data = resp.json()
    record = {**video, "language": data["language"], "transcript": data["transcript"]}
    out.write(json.dumps(record, ensure_ascii=False) + "\n")
    out.flush()
    counts["ok"] += 1


async def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("channel", help="@handle, channel URL or UC... ID")
    p.add_argument("--search", action="append", default=[], help="topic search inside the channel (repeatable)")
    p.add_argument("--out", default="transcripts.jsonl")
    p.add_argument("--csv", help="also write a CSV with one row per video")
    p.add_argument("--language", default="en")
    p.add_argument("--concurrency", type=int, default=5)
    p.add_argument("--dry-run", action="store_true", help="list videos and show the credit cost, fetch nothing")
    p.add_argument("--refresh", action="store_true", help="list the channel again instead of reusing the saved list")
    a = p.parse_args()

    key = os.environ.get("YTT_API_KEY")
    if not key:
        print("Set YTT_API_KEY first.", file=sys.stderr)
        return 1

    out_path = Path(a.out)
    done = set()
    if out_path.exists():
        with out_path.open(encoding="utf-8") as f:
            done = {json.loads(line)["video_id"] for line in f if line.strip()}

    # The video list is saved next to the output so a resumed run doesn't pay
    # for the listing again. --refresh lists the channel again.
    list_path = out_path.with_suffix(".videos.json")
    headers = {"Authorization": f"Bearer {key}"}
    rc = 0
    async with httpx.AsyncClient(headers=headers, timeout=60) as client:
        if list_path.exists() and not a.refresh:
            videos = json.loads(list_path.read_text(encoding="utf-8"))
            print(f"Using the saved list in {list_path} (no listing credits spent).")
        else:
            try:
                videos = await list_videos(client, a.channel, a.search)
            except OutOfCredits as e:
                print(f"Stopped: {e}", file=sys.stderr)
                return 2
            list_path.write_text(json.dumps(videos, ensure_ascii=False), encoding="utf-8")
            print(f"Listing cost {20 + len(a.search)} credits.")
        todo = [v for v in videos if v["video_id"] not in done]
        print(f"{len(videos)} videos, {len(done)} already saved, {len(todo)} to fetch (up to {len(todo)} credits).")
        if a.dry_run:
            return 0

        sem = asyncio.Semaphore(a.concurrency)
        counts = {"ok": 0, "no_captions": 0, "failed": 0}
        with out_path.open("a", encoding="utf-8") as out:
            tasks = [asyncio.create_task(fetch_one(client, sem, v, out, a.language, counts)) for v in todo]
            try:
                for finished in asyncio.as_completed(tasks):
                    await finished
            except OutOfCredits as e:
                print(f"Stopped: {e}. Top up and re-run; saved videos are skipped.", file=sys.stderr)
                rc = 2
                for t in tasks:
                    t.cancel()
                await asyncio.gather(*tasks, return_exceptions=True)
        print(
            f"Saved {counts['ok']} transcripts, {counts['no_captions']} had no captions, "
            f"{counts['failed']} failed."
        )

    if a.csv:
        with out_path.open(encoding="utf-8") as f, open(a.csv, "w", newline="", encoding="utf-8") as c:
            w = csv.writer(c)
            w.writerow(["video_id", "title", "language", "text"])
            for line in f:
                r = json.loads(line)
                text = " ".join(seg["text"] for seg in r["transcript"])
                w.writerow([r["video_id"], r["title"], r["language"], text])
        print(f"Wrote {a.csv}")
    return rc


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