"""Search everything a YouTube channel has said, with timestamped links.

    pip install httpx
    export YTT_API_KEY=tk_live_...
    python transcript_search.py index @TED --db ted.db --dry-run
    python transcript_search.py index @TED --db ted.db --search "climate"
    python transcript_search.py search "carbon tax" --db ted.db
    python transcript_search.py update @TED --db ted.db

Commands:
  index   Lists the channel's newest uploads, up to 100 (20 credits), plus
          optional topic searches inside the channel to reach older videos
          (1 credit each, up to 50 results). Fetches each transcript not yet
          in the database (1 credit per successful transcript) and indexes it.
          The video list is saved in the database, so a re-run after an
          interruption doesn't pay for it again; --refresh lists the channel again.
  update  Checks the channel's latest 15 uploads (free) and indexes only the
          new ones (1 credit each). Run it daily or weekly.
  import  Indexes a transcripts.jsonl file from channel_transcripts.py. No API
          calls, no credits.
  search  Full-text search of the local database. No API calls, no credits.
          Supports "exact phrases", prefix*, AND / OR / NOT and NEAR(a b, 10).

Transcripts are indexed in overlapping 30-second windows, so a phrase that
YouTube split across two caption lines is still found, and each link opens
the video at the caption line where the match starts. Failed calls, 404s
(no captions) and 429s are never charged.
"""

import argparse
import asyncio
import csv
import json
import os
import re
import sqlite3
import sys
import time

import httpx

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

SCHEMA = """
CREATE TABLE IF NOT EXISTS videos (
    video_id TEXT PRIMARY KEY,
    title TEXT,
    status TEXT NOT NULL,          -- 'ok' or 'no_captions'
    language TEXT,
    indexed_at INTEGER
);
CREATE VIRTUAL TABLE IF NOT EXISTS chunks USING fts5(
    text,
    video_id UNINDEXED,
    start UNINDEXED,
    tokenize = 'porter unicode61 remove_diacritics 2'
);
CREATE TABLE IF NOT EXISTS lines (video_id TEXT, start REAL, text TEXT);
CREATE INDEX IF NOT EXISTS lines_by_video ON lines (video_id, start);
CREATE TABLE IF NOT EXISTS listed (video_id TEXT PRIMARY KEY, title TEXT);
"""


class OutOfCredits(Exception):
    pass


def open_db(path: str) -> sqlite3.Connection:
    db = sqlite3.connect(path)
    db.executescript(SCHEMA)
    return db


def windows(segments: list[dict]) -> list[tuple[float, str]]:
    """Group caption segments into 30-second windows that start every 15 seconds."""
    out, next_start = [], 0.0
    for i, seg in enumerate(segments):
        start = float(seg.get("start", 0))
        if out and start < next_start:
            continue
        parts = []
        for later in segments[i:]:
            if float(later.get("start", 0)) >= start + WINDOW:
                break
            parts.append(later["text"].replace("\n", " "))
        out.append((start, " ".join(parts)))
        next_start = start + STEP
    return out


def save(db: sqlite3.Connection, video_id: str, title: str | None, language: str | None, segments: list[dict] | None):
    db.execute("DELETE FROM chunks WHERE video_id = ?", (video_id,))
    db.execute("DELETE FROM lines WHERE video_id = ?", (video_id,))
    if segments is None:
        status = "no_captions"
    else:
        status = "ok"
        db.executemany(
            "INSERT INTO chunks (text, video_id, start) VALUES (?, ?, ?)",
            [(text, video_id, start) for start, text in windows(segments)],
        )
        db.executemany(
            "INSERT INTO lines VALUES (?, ?, ?)",
            [(video_id, float(seg.get("start", 0)), seg["text"]) for seg in segments],
        )
    db.execute(
        "INSERT OR REPLACE INTO videos VALUES (?, ?, ?, ?, ?)",
        (video_id, title, status, language, int(time.time())),
    )
    db.commit()


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_channel(client: httpx.AsyncClient, channel: str, searches: list[str]) -> dict[str, str | None]:
    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"The channel has more than {len(videos)} uploads; use --search to reach older ones.")

    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 not False:
                videos.setdefault(v["videoId"], v.get("title"))
    return videos


async def list_latest(client: httpx.AsyncClient, channel: str) -> dict[str, str | None]:
    resp = await call(client, "/youtube/channel/latest", {"channel": channel})
    resp.raise_for_status()
    return {v["videoId"]: v.get("title") for v in resp.json()["results"]}


async def fetch_all(client, db, todo: dict[str, str | None], language: str, concurrency: int) -> int:
    sem = asyncio.Semaphore(concurrency)
    counts = {"ok": 0, "no_captions": 0, "failed": 0}

    async def one(video_id, title):
        async with sem:
            resp = await call(client, "/youtube/transcript", {"video_url": video_id, "language": language})
        if resp.status_code == 404:
            save(db, video_id, title, None, None)
            counts["no_captions"] += 1
        elif resp.status_code == 200:
            data = resp.json()
            save(db, video_id, title, data["language"], data["transcript"])
            counts["ok"] += 1
        else:
            print(f"{video_id}: HTTP {resp.status_code}, skipped (re-run to retry)", file=sys.stderr)
            counts["failed"] += 1

    rc = 0
    tasks = [asyncio.create_task(one(vid, title)) for vid, title in todo.items()]
    try:
        for finished in asyncio.as_completed(tasks):
            await finished
    except OutOfCredits as e:
        print(f"Stopped: {e}. Top up and re-run; indexed videos are skipped.", file=sys.stderr)
        rc = 2
        for t in tasks:
            t.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
    print(f"Indexed {counts['ok']} transcripts, {counts['no_captions']} had no captions, {counts['failed']} failed.")
    return rc


async def build(a) -> int:
    key = os.environ.get("YTT_API_KEY")
    if not key:
        print("Set YTT_API_KEY first.", file=sys.stderr)
        return 1
    db = open_db(a.db)
    known = {row[0] for row in db.execute("SELECT video_id FROM videos")}
    async with httpx.AsyncClient(headers={"Authorization": f"Bearer {key}"}, timeout=60) as client:
        try:
            listed = dict(db.execute("SELECT video_id, title FROM listed"))
            if a.command == "index" and listed and not a.refresh:
                videos = listed
                print(f"Using the channel list saved in {a.db} (no listing credits). --refresh lists it again.")
            elif a.command == "index":
                videos = await list_channel(client, a.channel, a.search)
                db.executemany("INSERT OR REPLACE INTO listed VALUES (?, ?)", videos.items())
                db.commit()
                print(f"Listing cost {20 + len(a.search)} credits.")
            else:
                videos = await list_latest(client, a.channel)
        except OutOfCredits as e:
            print(f"Stopped: {e}", file=sys.stderr)
            return 2
        todo = {vid: title for vid, title in videos.items() if vid not in known}
        print(f"{len(videos)} videos found, {len(videos) - len(todo)} already indexed, "
              f"{len(todo)} to fetch (up to {len(todo)} credits).")
        if a.dry_run or not todo:
            return 0
        return await fetch_all(client, db, todo, a.language, a.concurrency)


def import_jsonl(a) -> int:
    db = open_db(a.db)
    n = 0
    with open(a.file, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                r = json.loads(line)
                save(db, r["video_id"], r.get("title"), r.get("language"), r["transcript"])
                n += 1
    print(f"Imported {n} transcripts into {a.db}.")
    return 0


def clock(seconds: float) -> str:
    m, s = divmod(int(seconds), 60)
    h, m = divmod(m, 60)
    return f"{h}:{m:02d}:{s:02d}" if h else f"{m}:{s:02d}"


def exact_start(db: sqlite3.Connection, video_id: str, start: float, marked: str) -> float:
    """Find the caption line inside the window where the first matched word is spoken."""
    m = re.search("\x02(.*?)\x03", marked, re.DOTALL)
    if m and m.group(1).split():
        # A phrase match is marked as one span that may cross lines; its first word is enough.
        word = re.compile(r"\b" + re.escape(m.group(1).split()[0]) + r"\b", re.IGNORECASE)
        for line_start, text in db.execute(
            "SELECT start, text FROM lines WHERE video_id = ? AND start >= ? AND start < ? ORDER BY start",
            (video_id, start, start + WINDOW),
        ):
            if word.search(text):
                return line_start
    return start


def quoted(query: str) -> str:
    """Turn free text into a query FTS5 always accepts: every word becomes a quoted term."""
    return " ".join('"' + w.replace('"', '""') + '"' for w in query.split())


def search(a) -> int:
    db = open_db(a.db)
    sql = """
        SELECT videos.title, chunks.video_id, chunks.start,
               snippet(chunks, 0, '[', ']', '...', 16), highlight(chunks, 0, char(2), char(3))
        FROM chunks JOIN videos ON videos.video_id = chunks.video_id
        WHERE chunks MATCH ? ORDER BY rank LIMIT ?
    """
    try:
        rows = db.execute(sql, (a.query, a.limit * 3)).fetchall()
    except sqlite3.OperationalError:
        # Apostrophes, hyphens and colons are query syntax in FTS5. Retry as plain words.
        try:
            rows = db.execute(sql, (quoted(a.query), a.limit * 3)).fetchall()
        except sqlite3.OperationalError as e:
            print(f"Can't search for {a.query!r}: {e}", file=sys.stderr)
            return 1

    # Windows overlap, so one moment can match twice. Keep the better-ranked hit.
    hits, seen = [], []
    for title, video_id, start, snip, marked in rows:
        start = exact_start(db, video_id, start, marked)
        if any(v == video_id and abs(s - start) < WINDOW for v, s in seen):
            continue
        seen.append((video_id, start))
        hits.append((title, video_id, start, snip, f"https://youtu.be/{video_id}?t={int(start)}"))
        if len(hits) == a.limit:
            break

    for title, video_id, start, snip, link in hits:
        print(f"{title or video_id} @ {clock(start)}  {link}")
        print(f"    {snip}")
    indexed = db.execute("SELECT COUNT(*) FROM videos WHERE status = 'ok'").fetchone()[0]
    print(f"{len(hits)} matches in {indexed} indexed videos.")

    if a.csv:
        with open(a.csv, "w", newline="", encoding="utf-8") as c:
            w = csv.writer(c)
            w.writerow(["title", "video_id", "start", "snippet", "link"])
            w.writerows(hits)
        print(f"Wrote {a.csv}")
    return 0


def main() -> int:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="command", required=True)
    for name in ("index", "update"):
        s = sub.add_parser(name)
        s.add_argument("channel", help="@handle, channel URL or UC... ID")
        s.add_argument("--db", default="transcripts.db")
        s.add_argument("--language", default="en")
        s.add_argument("--concurrency", type=int, default=5)
        s.add_argument("--dry-run", action="store_true", help="show what would be fetched and the credit cost")
        if name == "index":
            s.add_argument("--search", action="append", default=[], help="topic search inside the channel (repeatable)")
            s.add_argument("--refresh", action="store_true", help="list the channel again instead of reusing the saved list")
    s = sub.add_parser("import")
    s.add_argument("file", help="transcripts.jsonl from channel_transcripts.py")
    s.add_argument("--db", default="transcripts.db")
    s = sub.add_parser("search")
    s.add_argument("query")
    s.add_argument("--db", default="transcripts.db")
    s.add_argument("--limit", type=int, default=20)
    s.add_argument("--csv", help="also write the matches to a CSV file")
    a = p.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")

    if a.command == "search":
        return search(a)
    if a.command == "import":
        return import_jsonl(a)
    return asyncio.run(build(a))


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