"""Download the transcript of every video in a YouTube playlist, in playlist order.

Writes one Markdown file per video, combined files sized for upload to
NotebookLM, Claude or ChatGPT, and a manifest.csv saying what happened to each
video. Re-running skips videos already saved; --refresh re-lists the playlist
to pick up new videos and renames files when the order has changed.

    pip install httpx
    export YTT_API_KEY=tk_live_...
    python playlist_transcripts.py "https://www.youtube.com/playlist?list=PL..." --dry-run
    python playlist_transcripts.py PL... --out course

Playlists longer than 100 videos: pass a free YouTube Data API key with
--youtube-key (or YOUTUBE_API_KEY) and the full list is read from Google,
50 videos per page.
"""
import argparse
import asyncio
import csv
import json
import os
import re
import sys
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import httpx

API = os.environ.get("YTT_API_BASE", "https://youtubetranscriptapi.com/api/v2")
GOOGLE = "https://www.googleapis.com/youtube/v3/playlistItems"
RETRY_STATUSES = {408, 429, 500, 502, 503}


class CreditsExhausted(Exception):
    pass


def playlist_id(value):
    """Accept a playlist URL or a bare ID."""
    if "list=" in value:
        return parse_qs(urlparse(value).query)["list"][0]
    return value.strip()


def slug(title, limit=60):
    clean = re.sub(r"[^A-Za-z0-9]+", "-", title or "untitled").strip("-")
    return clean[:limit] or "untitled"


async def get(client, path, params, attempts=5):
    """GET with retries on rate limits and transient errors. 402 stops the run."""
    for attempt in range(attempts):
        try:
            resp = await client.get(path, params=params)
        except httpx.TransportError:
            await asyncio.sleep(2 ** attempt)
            continue
        if resp.status_code == 402:
            raise CreditsExhausted(resp.json().get("detail", "out of credits"))
        if resp.status_code not in RETRY_STATUSES:
            return resp
        try:
            delay = float(resp.headers.get("Retry-After", ""))
        except ValueError:
            delay = 2 ** attempt
        await asyncio.sleep(min(delay, 60))
    return resp


async def list_with_api(client, pid):
    """First page of the playlist from YouTubeTranscriptAPI: up to 100 videos, 20 credits."""
    resp = await get(client, "/youtube/playlist/videos", {"playlist": pid})
    resp.raise_for_status()
    body = resp.json()
    videos = [{"position": int(v.get("index", i)) + 1, "id": v["videoId"], "title": v.get("title") or ""}
              for i, v in enumerate(body["results"])]
    return videos, body.get("playlist_info") or {}, bool(body.get("has_more"))


async def list_with_google(pid, key):
    """Every item in the playlist from the YouTube Data API: 1 quota unit per 50 videos."""
    videos, token = [], None
    async with httpx.AsyncClient(timeout=30) as google:
        while True:
            params = {"part": "snippet", "playlistId": pid, "maxResults": 50, "key": key}
            if token:
                params["pageToken"] = token
            resp = await google.get(GOOGLE, params=params)
            resp.raise_for_status()
            page = resp.json()
            for item in page.get("items", []):
                s = item["snippet"]
                videos.append({"position": s["position"] + 1, "id": s["resourceId"]["videoId"], "title": s.get("title") or ""})
            token = page.get("nextPageToken")
            if not token:
                return videos


def file_for(folder, video):
    return folder / "videos" / f"{video['position']:03d}-{slug(video['title'], 40)}-{video['id']}.md"


def saved_file(folder, video):
    """The saved file for this video, whatever position it was saved under."""
    return next((folder / "videos").glob(f"*-{video['id']}.md"), None)


def to_markdown(video, data):
    lines = [f"## {video['position']:03d}. {video['title']}", "",
             f"https://www.youtube.com/watch?v={video['id']}", ""]
    for seg in data["transcript"]:
        m, s = divmod(int(seg["start"]), 60)
        lines.append(f"[{m}:{s:02d}] {seg['text']}")
    return "\n".join(lines) + "\n"


async def fetch_one(client, video, folder, language, gate, status, stop):
    path = file_for(folder, video)
    old = saved_file(folder, video)
    if old:
        if old != path:  # the playlist was reordered: rename for free
            text = old.read_text(encoding="utf-8").split("\n", 1)[1]
            path.write_text(f"## {video['position']:03d}. {video['title']}\n{text}", encoding="utf-8")
            old.unlink()
        status[video["id"]] = "saved earlier"
        return
    async with gate:
        if stop.is_set():
            return
        try:
            resp = await get(client, "/youtube/transcript", {"video_url": video["id"], "language": language})
        except CreditsExhausted as e:
            if not stop.is_set():
                print(f"Stopped: {e}. Everything fetched so far is saved; top up and run again.")
            stop.set()
            return
    if resp.status_code == 200:
        data = resp.json()
        path.write_text(to_markdown(video, data), encoding="utf-8")
        got = data.get("language", language)
        status[video["id"]] = "saved" if got == language else f"saved (language: {got})"
        print(f"  {video['position']:03d} saved  {video['title'][:60]}")
    elif resp.status_code == 404:
        status[video["id"]] = "no captions, private or removed"
    else:
        status[video["id"]] = f"failed ({resp.status_code})"


def combine(folder, videos, max_words):
    """Join the saved files, in playlist order, into parts of at most max_words words."""
    parts, current, words = [], [], 0
    for f in filter(None, (saved_file(folder, v) for v in sorted(videos, key=lambda v: v["position"]))):
        text = f.read_text(encoding="utf-8")
        n = len(text.split())
        if current and words + n > max_words:
            parts.append(current)
            current, words = [], 0
        current.append(text)
        words += n
    if current:
        parts.append(current)
    for old in folder.glob("playlist-part-*.md"):
        old.unlink()
    for i, chunk in enumerate(parts, 1):
        (folder / f"playlist-part-{i}.md").write_text("\n".join(chunk), encoding="utf-8")
    return len(parts)


async def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("playlist", help="playlist URL or ID")
    ap.add_argument("--out", help="output folder (default: the playlist ID)")
    ap.add_argument("--language", default="en")
    ap.add_argument("--max-words", type=int, default=100_000, help="words per combined file")
    ap.add_argument("--youtube-key", default=os.environ.get("YOUTUBE_API_KEY"),
                    help="YouTube Data API key, for playlists over 100 videos")
    ap.add_argument("--refresh", action="store_true", help="list the playlist again to pick up changes")
    ap.add_argument("--dry-run", action="store_true", help="list the playlist and print the cost, fetch nothing")
    args = ap.parse_args()

    pid = playlist_id(args.playlist)
    folder = Path(args.out or pid)
    (folder / "videos").mkdir(parents=True, exist_ok=True)
    cache = folder / ".videos.json"
    headers = {"Authorization": f"Bearer {os.environ['YTT_API_KEY']}"}

    async with httpx.AsyncClient(base_url=API, headers=headers, timeout=60) as client:
        try:
            if cache.exists() and not args.refresh:
                videos = json.loads(cache.read_text(encoding="utf-8"))
            elif args.youtube_key:
                videos = await list_with_google(pid, args.youtube_key)
            else:
                videos, info, more = await list_with_api(client, pid)
                print(f"{info.get('title', pid)}: {info.get('numVideos', '?')}")
                if more:
                    print("Note: this playlist has more than 100 videos. Only the first 100 are listed;"
                          " pass --youtube-key to get all of them.")
            cache.write_text(json.dumps(videos), encoding="utf-8")

            todo = [v for v in videos if not saved_file(folder, v)]
            print(f"{len(videos)} videos listed, {len(todo)} to fetch, up to {len(todo)} transcript credits.")
            if args.dry_run:
                return 0

            status, gate, stop = {}, asyncio.Semaphore(5), asyncio.Event()
            await asyncio.gather(*(fetch_one(client, v, folder, args.language, gate, status, stop) for v in videos))
            exit_code = 2 if stop.is_set() else 0
        except CreditsExhausted as e:
            print(f"Stopped before listing: {e}")
            return 2

    with open(folder / "manifest.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["position", "video_id", "title", "status"])
        for v in videos:
            saved = saved_file(folder, v) is not None
            w.writerow([v["position"], v["id"], v["title"], status.get(v["id"], "saved earlier" if saved else "not fetched")])
    parts = combine(folder, videos, args.max_words)
    done = sum(1 for s in status.values() if s.startswith("saved"))
    print(f"{done} of {len(videos)} transcripts on disk, in {parts} combined file(s). See manifest.csv.")
    return exit_code


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