"""Turn YouTube transcripts into SRT and WebVTT subtitle files.

    pip install requests
    export YTT_API_KEY=tk_live_...
    python transcript_to_subtitles.py jNQXAC9IVRw
    python transcript_to_subtitles.py --playlist PLxxxxxxxx --format both
    python transcript_to_subtitles.py --channel @TED --max-videos 20 --language es

For each video it fetches the transcript with timestamps (1 credit; videos
without captions return 404 and cost nothing), then:
  1. Cleans the text and fixes overlapping timings (automatic captions often
     run past the start of the next line, which doubles lines in editors).
  2. Regroups the short caption fragments into readable cues: at most
     --max-lines lines of --max-chars characters, and no longer than
     --max-duration seconds on screen.
  3. Writes OUT/VIDEO_ID.LANG.srt and/or .vtt in UTF-8.

Videos that already have a file in OUT are skipped, so a re-run after an
error or a 402 only pays for the videos still missing. (If a video has no
track in --language, the API returns another one and the file is named after
it. To fetch a second language for the same videos, use another --out.)
Use --raw to keep the original caption lines exactly as they came.
"""

import argparse
import html
import os
import re
import sys
import time
from pathlib import Path

import requests

API = os.environ.get("YTT_API_BASE", "https://youtubetranscriptapi.com/api/v2")
RETRYABLE = {408, 429, 500, 503}
SOUND_TAG = re.compile(r"^\s*[\[(][^\])]{1,40}[\])]\s*$")  # a whole line like [Music] or (applause)


class OutOfCredits(Exception):
    pass


def call(session, path, params, tries=4):
    """GET with retries on transient errors, honouring Retry-After."""
    for attempt in range(tries):
        resp = session.get(f"{API}{path}", params=params, timeout=60)
        if resp.status_code == 402:
            detail = resp.json().get("detail")
            raise OutOfCredits(detail.get("message") if isinstance(detail, dict) else detail or "out of credits")
        if resp.status_code not in RETRYABLE:
            return resp
        time.sleep(min(float(resp.headers.get("Retry-After", 2 ** attempt)), 60))
    return resp


# ---------- timing and text ----------

def clean_segments(raw, drop_sound_tags):
    """Floats for times, tidy text, sorted by start, empty and duplicate lines removed."""
    segs = []
    for s in raw:
        text = " ".join(html.unescape(str(s.get("text", ""))).split())
        if not text or (drop_sound_tags and SOUND_TAG.match(text)):
            continue
        start = float(s.get("start", 0))
        end = start + max(float(s.get("duration", 0)), 0.0)
        segs.append({"start": start, "end": end, "text": text})
    segs.sort(key=lambda s: s["start"])
    out = []
    for s in segs:
        if out and s["text"] == out[-1]["text"] and s["start"] <= out[-1]["end"] + 0.05:
            out[-1]["end"] = max(out[-1]["end"], s["end"])  # the same line repeated by rolling captions
            continue
        out.append(s)
    return out


def fix_overlaps(segs):
    """End each line no later than the next one starts, and never before it begins."""
    out = []
    for i, cur in enumerate(segs):
        nxt = segs[i + 1] if i + 1 < len(segs) else None
        if nxt:
            cur["end"] = min(cur["end"], nxt["start"])
            if cur["end"] - cur["start"] < 0.05:  # starts with the next line: fold its text in, don't lose it
                nxt["text"] = f"{cur['text']} {nxt['text']}"
                continue
        cur["end"] = max(cur["end"], cur["start"] + 0.05)
        out.append(cur)
    return out


def split_long(seg, limit):
    """Split one over-long caption line into pieces of at most `limit` chars, sharing its time by length."""
    words, pieces, cur = seg["text"].split(), [], ""
    for w in words:
        if cur and len(cur) + 1 + len(w) > limit:
            pieces.append(cur)
            cur = w
        else:
            cur = f"{cur} {w}".strip()
    pieces.append(cur)
    total = sum(len(p) for p in pieces) or 1
    t, span, out = seg["start"], seg["end"] - seg["start"], []
    for p in pieces:
        d = span * len(p) / total
        out.append({"start": t, "end": t + d, "text": p})
        t += d
    return out


def split_slow(seg, max_duration):
    """Split a line that stays up longer than max_duration into word groups that share its time."""
    words = seg["text"].split()
    n = min(int((seg["end"] - seg["start"]) // max_duration) + 1, len(words))
    if n < 2:
        return [seg]
    size = -(-len(words) // n)
    longest = max(len(" ".join(words[i:i + size])) for i in range(0, len(words), size))
    return split_long(seg, longest)


def regroup(segs, max_chars, max_lines, max_duration, max_gap=1.0):
    """Merge short fragments into cues that fit the screen and the time limit."""
    limit = max_chars * max_lines
    pieces = [p for s in segs for p in (split_long(s, limit) if len(s["text"]) > limit else [s])]
    pieces = [q for p in pieces for q in (split_slow(p, max_duration) if p["end"] - p["start"] > max_duration else [p])]
    cues = []
    for p in pieces:
        if cues and not (SOUND_TAG.match(p["text"]) or SOUND_TAG.match(cues[-1]["text"])):  # keep [Music] on its own
            c = cues[-1]
            fits = len(c["text"]) + 1 + len(p["text"]) <= limit and wrap(f"{c['text']} {p['text']}", max_chars, max_lines)
            if fits and p["end"] - c["start"] <= max_duration and p["start"] - c["end"] <= max_gap:
                c["text"] = f"{c['text']} {p['text']}"
                c["end"] = p["end"]
                continue
        cues.append(dict(p))
    return cues


def apply_min_duration(cues, min_duration):
    """Let very short cues stay up longer, but never into the next cue."""
    for i, c in enumerate(cues):
        if c["end"] - c["start"] < min_duration:
            ceiling = cues[i + 1]["start"] if i + 1 < len(cues) else c["start"] + min_duration
            c["end"] = max(c["end"], min(c["start"] + min_duration, ceiling))
    return cues


def wrap(text, max_chars, max_lines):
    """Break text into at most max_lines lines of max_chars, as evenly as possible. None if it can't fit."""
    if len(text) <= max_chars:
        return [text]
    if max_lines < 2:
        return None
    words = text.split()
    best = None
    for i in range(1, len(words)):
        a, b = " ".join(words[:i]), " ".join(words[i:])
        if len(a) <= max_chars and len(b) <= max_chars:
            score = abs(len(a) - len(b)) - (6 if a[-1] in ",.;:?!" else 0)  # prefer breaking after punctuation
            if best is None or score < best[0]:
                best = (score, [a, b])
    return best[1] if best else None


def build_cues(raw, a):
    segs = fix_overlaps(clean_segments(raw, a.drop_sound_tags))
    if a.raw:
        return [dict(s, lines=[s["text"]]) for s in segs]
    cues = apply_min_duration(regroup(segs, a.max_chars, a.max_lines, a.max_duration), a.min_duration)
    for c in cues:
        c["lines"] = wrap(c["text"], a.max_chars, a.max_lines) or [c["text"]]
    return cues


# ---------- writers ----------

def stamp(seconds, sep):
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3_600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d}{sep}{ms:03d}"


def to_srt(cues):
    blocks = [f"{i}\n{stamp(c['start'], ',')} --> {stamp(c['end'], ',')}\n" + "\n".join(c["lines"]).replace("-->", "->")
              for i, c in enumerate(cues, 1)]
    return "\n\n".join(blocks) + "\n"


def to_vtt(cues):
    blocks = [f"{stamp(c['start'], '.')} --> {stamp(c['end'], '.')}\n" + "\n".join(c["lines"]).replace("-->", "->")
              for c in cues]
    return "WEBVTT\n\n" + "\n\n".join(blocks) + "\n"


def too_fast(cues, cps):
    return sum(1 for c in cues if len(c["text"]) / max(c["end"] - c["start"], 0.001) > cps)


# ---------- main ----------

def glob_escape(name):
    return re.sub(r"([\[\]*?])", r"[\1]", name)


def list_videos(session, a):
    ids = list(a.videos)
    for kind, path, key in (("playlist", "/youtube/playlist/videos", "playlist"), ("channel", "/youtube/channel/videos", "channel")):
        value = getattr(a, kind)
        if value:
            resp = call(session, path, {key: value})
            if resp.status_code != 200:
                sys.exit(f"{kind} {value}: HTTP {resp.status_code} {resp.text[:200]}")
            found = [v["videoId"] for v in resp.json().get("results", [])]
            print(f"{kind} {value}: {len(found)} videos (20 credits).")
            ids += found
    seen = set()
    ids = [v for v in ids if not (v in seen or seen.add(v))]
    return ids[: a.max_videos] if a.max_videos else ids


def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("videos", nargs="*", help="video IDs or youtube.com / youtu.be URLs")
    p.add_argument("--playlist", help="playlist ID or URL (lists up to 100 videos, 20 credits)")
    p.add_argument("--channel", help="@handle, channel URL or UC... ID (newest 100 uploads, 20 credits)")
    p.add_argument("--max-videos", type=int, help="stop after this many videos")
    p.add_argument("--format", choices=["srt", "vtt", "both"], default="srt")
    p.add_argument("--language", default="en", help="preferred caption language, e.g. en, es, de")
    p.add_argument("--out", default="subtitles")
    p.add_argument("--max-chars", type=int, default=42, help="characters per line")
    p.add_argument("--max-lines", type=int, default=2, help="lines per cue")
    p.add_argument("--max-duration", type=float, default=7.0, help="longest time a cue stays on screen, seconds")
    p.add_argument("--min-duration", type=float, default=5 / 6, help="shortest cue, seconds (if the next cue allows)")
    p.add_argument("--cps", type=float, default=20.0, help="reading speed to warn about, characters per second")
    p.add_argument("--drop-sound-tags", action="store_true", help="remove lines like [Music] or [Applause]")
    p.add_argument("--raw", action="store_true", help="keep the original caption lines, only fix overlaps")
    a = p.parse_args()
    sys.stdout.reconfigure(encoding="utf-8", line_buffering=True)

    key = os.environ.get("YTT_API_KEY")
    if not key or not (a.videos or a.playlist or a.channel):
        p.error("set YTT_API_KEY and pass at least one video, --playlist or --channel")
    formats = ["srt", "vtt"] if a.format == "both" else [a.format]
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)

    session = requests.Session()
    session.headers["Authorization"] = f"Bearer {key}"
    stats = {"written": 0, "skipped": 0, "no_captions": 0, "failed": 0, "credits": 0}
    rc = 0
    try:
        videos = list_videos(session, a)
        print(f"{len(videos)} video(s): up to {len(videos)} credit(s) for transcripts.")
        for v in videos:
            vid = re.search(r"(?:v=|youtu\.be/|shorts/|^)([\w-]{11})(?![\w-])", v)
            vid = vid.group(1) if vid else v
            if all(any(out.glob(f"{glob_escape(vid)}.*.{f}")) for f in formats):
                stats["skipped"] += 1  # already converted (in whatever language came back)
                continue
            resp = call(session, "/youtube/transcript", {"video_url": v, "language": a.language, "format": "json",
                                                         "include_timestamp": "true", "send_metadata": "true"})
            if resp.status_code == 404:
                print(f"  {vid}: no captions (not charged)")
                stats["no_captions"] += 1
                continue
            if resp.status_code != 200:
                print(f"  {vid}: HTTP {resp.status_code}, skipped", file=sys.stderr)
                stats["failed"] += 1
                continue
            stats["credits"] += int(resp.headers.get("X-Credits-Charged", 1))
            data = resp.json()
            cues = build_cues(data.get("transcript") or [], a)
            if not cues:
                print(f"  {vid}: transcript is empty, nothing written")
                stats["failed"] += 1
                continue
            lang = re.sub(r"[^\w-]", "", data.get("language") or a.language) or a.language
            if lang != a.language:
                print(f"  {vid}: no '{a.language}' captions, saved the '{lang}' track instead")
            targets = {f: out / f"{vid}.{lang}.{f}" for f in formats}
            for f, path in targets.items():
                path.write_text(to_srt(cues) if f == "srt" else to_vtt(cues), encoding="utf-8")
            title = (data.get("metadata") or {}).get("title", "")
            fast = too_fast(cues, a.cps)
            note = f", {fast} cue(s) faster than {a.cps:g} chars/s" if fast else ""
            print(f"  {vid}: {len(cues)} cues -> {', '.join(t.name for t in targets.values())}{note}  {title}")
            stats["written"] += 1
    except OutOfCredits as e:
        print(f"Stopped: {str(e).rstrip('.')}. Run again after topping up; finished files are skipped.", file=sys.stderr)
        rc = 2
    print("Done: {written} written, {skipped} already there, {no_captions} without captions, "
          "{failed} failed, {credits} credit(s) used.".format(**stats))
    return rc


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