youtube_transcript_tool.py

#!/usr/bin/env python3
"""Local YouTube transcript toolkit (no MCP).

Modes:
- transcript: fetch transcript text or timestamped entries
- languages: list available transcript languages
- analyze: extractive summary, key points, or quotes
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from collections import Counter
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Sequence, Tuple
from urllib.parse import parse_qs, urlparse

try:
    from youtube_transcript_api import YouTubeTranscriptApi
except Exception as import_error:  # pragma: no cover - import is environment-dependent
    YouTubeTranscriptApi = None
    IMPORT_ERROR = import_error
else:
    IMPORT_ERROR = None

VIDEO_ID_RE = re.compile(r"^[A-Za-z0-9_-]{11}$")
SENTENCE_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")
WORD_RE = re.compile(r"[A-Za-z0-9']+")

STOPWORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "for",
    "from",
    "has",
    "he",
    "in",
    "is",
    "it",
    "its",
    "of",
    "on",
    "or",
    "that",
    "the",
    "to",
    "was",
    "were",
    "will",
    "with",
    "you",
    "your",
    "this",
    "they",
    "we",
    "i",
    "our",
}


@dataclass
class ToolError(Exception):
    code: str
    message: str
    hint: str


def parse_bool(value: str) -> bool:
    lowered = value.strip().lower()
    if lowered in {"true", "1", "yes", "y"}:
        return True
    if lowered in {"false", "0", "no", "n"}:
        return False
    raise argparse.ArgumentTypeError(f"Invalid boolean value: {value}")


def normalize_video_id(url_or_id: str) -> str:
    candidate = url_or_id.strip()
    if VIDEO_ID_RE.match(candidate):
        return candidate

    parsed = urlparse(candidate)
    host = parsed.netloc.lower().replace("www.", "").replace("m.", "")
    path = parsed.path.strip("/")
    video_id = None

    if host == "youtu.be" and path:
        video_id = path.split("/")[0]
    elif "youtube.com" in host:
        if path == "watch":
            video_id = parse_qs(parsed.query).get("v", [None])[0]
        else:
            parts = [p for p in path.split("/") if p]
            if len(parts) >= 2 and parts[0] in {"embed", "v", "shorts", "live", "e"}:
                video_id = parts[1]
            elif parts and VIDEO_ID_RE.match(parts[0]):
                video_id = parts[0]

    if not video_id and (not host or host == "youtu.be" or "youtube.com" in host):
        match = re.search(r"(?:v=|/)([A-Za-z0-9_-]{11})(?:[?&/#]|$)", candidate)
        if match:
            video_id = match.group(1)

    if video_id and VIDEO_ID_RE.match(video_id):
        return video_id

    raise ToolError(
        "invalid_video_id",
        "Could not extract a valid YouTube video ID from the provided input.",
        "Use a full YouTube URL or a raw 11-character ID like dQw4w9WgXcQ.",
    )


def ensure_dependency() -> None:
    if YouTubeTranscriptApi is None:
        detail = str(IMPORT_ERROR) if IMPORT_ERROR else "youtube-transcript-api missing"
        raise ToolError(
            "dependency_missing",
            f"youtube-transcript-api is not available: {detail}",
            "Run: python3 -m pip install youtube-transcript-api",
        )


def map_exception(error: Exception) -> ToolError:
    name = error.__class__.__name__.lower()
    message = str(error)
    lowered = message.lower()

    if "invalidvideoid" in name or ("invalid" in lowered and "video" in lowered and "id" in lowered):
        return ToolError(
            "invalid_video_id",
            message,
            "Check your URL format or pass a raw 11-character YouTube video ID.",
        )

    if "transcriptsdisabled" in name or "disabled" in lowered:
        return ToolError(
            "transcripts_disabled",
            message,
            "Captions are disabled for this video. Try a different video.",
        )

    if "notranscriptfound" in name or "no transcript" in lowered:
        return ToolError(
            "no_transcript",
            message,
            "Try another language or use --mode languages to inspect available transcripts.",
        )

    if "videounavailable" in name or ("video" in lowered and "unavailable" in lowered):
        return ToolError(
            "video_unavailable",
            message,
            "The video may be private, deleted, or region/age restricted.",
        )

    if "toomanyrequests" in name or "429" in lowered or "rate" in lowered:
        return ToolError(
            "rate_limited",
            message,
            "Retry later or from a different network.",
        )

    if "couldnotretrievetranscript" in name:
        return ToolError(
            "transcript_fetch_failed",
            message,
            "Retry once, then try --mode languages to inspect availability.",
        )

    return ToolError(
        "transcript_fetch_failed",
        message or "Failed to retrieve transcript data.",
        "Retry with a different video or language.",
    )


def iter_transcript_objects(transcript_list_obj: Any) -> Iterable[Any]:
    try:
        for item in transcript_list_obj:
            yield item
        return
    except TypeError:
        pass

    for attr in ("transcripts", "_manually_created_transcripts", "_generated_transcripts"):
        if not hasattr(transcript_list_obj, attr):
            continue

        value = getattr(transcript_list_obj, attr)
        if isinstance(value, dict):
            for item in value.values():
                yield item
        elif isinstance(value, list):
            for item in value:
                yield item


def list_languages(video_id: str) -> List[Dict[str, Any]]:
    ensure_dependency()

    try:
        api = YouTubeTranscriptApi() if callable(YouTubeTranscriptApi) else None

        if hasattr(YouTubeTranscriptApi, "list_transcripts"):
            transcript_list_obj = YouTubeTranscriptApi.list_transcripts(video_id)
        elif api is not None and hasattr(api, "list_transcripts"):
            transcript_list_obj = api.list_transcripts(video_id)
        elif api is not None and hasattr(api, "list"):
            transcript_list_obj = api.list(video_id)
        else:
            raise ToolError(
                "api_compat_error",
                "Unable to find a compatible list_transcripts() API.",
                "Upgrade youtube-transcript-api to the latest version.",
            )

        languages = []
        for transcript in iter_transcript_objects(transcript_list_obj):
            code = getattr(transcript, "language_code", None)
            name = getattr(transcript, "language", None)
            is_generated = bool(getattr(transcript, "is_generated", False))

            if not code:
                continue

            languages.append(
                {
                    "language_code": code,
                    "language_name": name or code,
                    "is_generated": is_generated,
                }
            )

        if not languages:
            raise ToolError(
                "no_transcript",
                "No transcript languages were found for this video.",
                "The video may not have captions enabled.",
            )

        # Deterministic order: human captions first, then generated; then by code.
        languages.sort(key=lambda x: (x["is_generated"], x["language_code"]))
        return languages

    except ToolError:
        raise
    except Exception as error:  # pragma: no cover - depends on third-party exceptions
        raise map_exception(error) from error


def normalize_entries(raw_entries: Any) -> List[Dict[str, Any]]:
    if hasattr(raw_entries, "to_raw_data"):
        raw_entries = raw_entries.to_raw_data()

    if not isinstance(raw_entries, list):
        raise ToolError(
            "transcript_parse_error",
            "Transcript API returned an unsupported data shape.",
            "Update youtube-transcript-api and retry.",
        )

    entries: List[Dict[str, Any]] = []
    for item in raw_entries:
        if isinstance(item, dict):
            text = str(item.get("text", "")).strip()
            start = float(item.get("start", 0.0) or 0.0)
            duration = float(item.get("duration", 0.0) or 0.0)
        else:
            text = str(getattr(item, "text", "")).strip()
            start = float(getattr(item, "start", 0.0) or 0.0)
            duration = float(getattr(item, "duration", 0.0) or 0.0)

        if not text:
            continue

        entries.append({"text": text, "start": start, "duration": duration})

    if not entries:
        raise ToolError(
            "no_transcript",
            "Transcript payload was empty.",
            "Try another language or a different video.",
        )

    return entries


def fetch_entries(video_id: str, language_code: str) -> List[Dict[str, Any]]:
    ensure_dependency()

    api = YouTubeTranscriptApi() if callable(YouTubeTranscriptApi) else None
    calls = []

    if hasattr(YouTubeTranscriptApi, "get_transcript"):
        calls.append(lambda: YouTubeTranscriptApi.get_transcript(video_id, languages=[language_code]))
    if api is not None and hasattr(api, "get_transcript"):
        calls.append(lambda: api.get_transcript(video_id, languages=[language_code]))
    if api is not None and hasattr(api, "fetch"):
        calls.append(lambda: api.fetch(video_id, languages=[language_code]))
    if hasattr(YouTubeTranscriptApi, "fetch"):
        calls.append(lambda: YouTubeTranscriptApi.fetch(video_id, languages=[language_code]))

    if not calls:
        raise ToolError(
            "api_compat_error",
            "Unable to find a compatible transcript fetch API.",
            "Upgrade youtube-transcript-api to the latest version.",
        )

    last_error: Exception | None = None
    for call in calls:
        try:
            return normalize_entries(call())
        except ToolError:
            raise
        except Exception as error:  # pragma: no cover - depends on third-party exceptions
            last_error = error

    if last_error is not None:
        raise map_exception(last_error) from last_error

    raise ToolError(
        "transcript_fetch_failed",
        "Transcript fetch failed for an unknown reason.",
        "Retry with a different language or video.",
    )


def build_plain_text(entries: Sequence[Dict[str, Any]]) -> str:
    return " ".join(entry["text"].strip() for entry in entries if entry.get("text")).strip()


def split_sentences(text: str) -> List[str]:
    if not text:
        return []
    sentences = [s.strip() for s in SENTENCE_SPLIT_RE.split(text) if s.strip()]
    return sentences


def tokenize(text: str) -> List[str]:
    return [w.lower() for w in WORD_RE.findall(text)]


def sentence_scores(sentences: Sequence[str]) -> List[Tuple[float, int, str]]:
    corpus_tokens = [w for sentence in sentences for w in tokenize(sentence) if w not in STOPWORDS and len(w) > 2]
    if not corpus_tokens:
        return [(1.0, idx, sentence) for idx, sentence in enumerate(sentences)]

    frequencies = Counter(corpus_tokens)
    scored = []
    for idx, sentence in enumerate(sentences):
        tokens = [w for w in tokenize(sentence) if w not in STOPWORDS and len(w) > 2]
        if not tokens:
            scored.append((0.0, idx, sentence))
            continue
        unique_tokens = set(tokens)
        score = sum(frequencies[token] for token in unique_tokens) / max(len(tokens), 1)
        scored.append((score, idx, sentence))
    return scored


def analyze_summary(entries: Sequence[Dict[str, Any]], max_items: int) -> Dict[str, Any]:
    text = build_plain_text(entries)
    sentences = split_sentences(text)

    if not sentences:
        return {
            "analysis_type": "summary",
            "summary": "No usable transcript text was found.",
            "key_takeaways": [],
        }

    ranked = sorted(sentence_scores(sentences), key=lambda x: x[0], reverse=True)
    pick_count = min(max(3, min(max_items, 5)), len(sentences))
    selected_idxs = sorted(idx for _, idx, _ in ranked[:pick_count])
    selected = [sentences[idx] for idx in selected_idxs]

    summary = " ".join(selected[: min(3, len(selected))]).strip()

    return {
        "analysis_type": "summary",
        "summary": summary,
        "key_takeaways": selected[:max_items],
    }


def analyze_key_points(entries: Sequence[Dict[str, Any]], max_items: int) -> Dict[str, Any]:
    text = build_plain_text(entries)
    sentences = split_sentences(text)

    if not sentences:
        return {"analysis_type": "key_points", "key_points": []}

    ranked = sorted(sentence_scores(sentences), key=lambda x: x[0], reverse=True)
    top = ranked[: min(max_items, len(ranked))]

    return {
        "analysis_type": "key_points",
        "key_points": [sentence for _, _, sentence in top],
    }


def analyze_quotes(entries: Sequence[Dict[str, Any]], max_items: int) -> Dict[str, Any]:
    candidates = []
    seen = set()

    for entry in entries:
        text = entry["text"].strip()
        if len(tokenize(text)) < 6:
            continue
        normalized = text.lower()
        if normalized in seen:
            continue
        seen.add(normalized)
        candidates.append((len(text), text, entry["start"]))

    if not candidates:
        return {"analysis_type": "quotes", "quotes": []}

    candidates.sort(key=lambda x: x[0], reverse=True)
    selected = candidates[:max_items]

    return {
        "analysis_type": "quotes",
        "quotes": [
            {
                "text": text,
                "start": start,
            }
            for _, text, start in selected
        ],
    }


def format_timestamp(seconds: float) -> str:
    whole = int(seconds)
    hours = whole // 3600
    minutes = (whole % 3600) // 60
    secs = whole % 60
    if hours > 0:
        return f"{hours}:{minutes:02d}:{secs:02d}"
    return f"{minutes}:{secs:02d}"


def emit(data: Dict[str, Any], as_json: bool) -> None:
    if as_json:
        print(json.dumps(data, ensure_ascii=False))
        return

    if not data.get("ok"):
        error = data["error"]
        print(f"Error [{error['code']}]: {error['message']}")
        print(f"Hint: {error['hint']}")
        return

    mode = data.get("mode")
    print(f"Mode: {mode}")
    print(f"Video ID: {data.get('video_id')}")

    if mode == "languages":
        print("Languages:")
        for item in data.get("languages", []):
            suffix = " (auto-generated)" if item.get("is_generated") else ""
            print(f"- {item['language_code']}: {item['language_name']}{suffix}")

    if mode == "transcript":
        transcript = data.get("transcript")
        if isinstance(transcript, str):
            print("\nTranscript:\n")
            print(transcript)
        else:
            print("\nTranscript entries:\n")
            for entry in transcript:
                print(f"[{format_timestamp(entry['start'])}] {entry['text']}")

    if mode == "analyze":
        analysis = data.get("analysis", {})
        print("\nAnalysis:\n")
        print(json.dumps(analysis, ensure_ascii=False, indent=2))


def success_payload(mode: str, video_id: str, **kwargs: Any) -> Dict[str, Any]:
    payload = {"ok": True, "mode": mode, "video_id": video_id}
    payload.update(kwargs)
    return payload


def error_payload(error: ToolError) -> Dict[str, Any]:
    return {
        "ok": False,
        "error": {
            "code": error.code,
            "message": error.message,
            "hint": error.hint,
        },
    }


def run(args: argparse.Namespace) -> Dict[str, Any]:
    video_id = normalize_video_id(args.url)

    if args.mode == "languages":
        languages = list_languages(video_id)
        return success_payload(
            "languages",
            video_id,
            count=len(languages),
            languages=languages,
        )

    languages = list_languages(video_id)
    available_codes = [item["language_code"] for item in languages]
    preferred = args.lang.strip() if args.lang else "en"
    language_used = preferred if preferred in available_codes else available_codes[0]

    entries = fetch_entries(video_id, language_used)

    if args.mode == "transcript":
        transcript_data: Any
        if args.include_timestamps:
            transcript_data = entries
        else:
            transcript_data = build_plain_text(entries)

        return success_payload(
            "transcript",
            video_id,
            language_requested=preferred,
            language_used=language_used,
            include_timestamps=args.include_timestamps,
            entry_count=len(entries),
            transcript=transcript_data,
        )

    if args.analysis_type == "summary":
        analysis = analyze_summary(entries, args.max_items)
    elif args.analysis_type == "key_points":
        analysis = analyze_key_points(entries, args.max_items)
    else:
        analysis = analyze_quotes(entries, args.max_items)

    return success_payload(
        "analyze",
        video_id,
        language_requested=preferred,
        language_used=language_used,
        entry_count=len(entries),
        analysis=analysis,
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="YouTube transcript toolkit (no MCP)")
    parser.add_argument("--mode", choices=["transcript", "languages", "analyze"], required=True)
    parser.add_argument("--url", required=True, help="YouTube URL or 11-character video ID")
    parser.add_argument("--lang", default="en", help="Preferred language code (default: en)")
    parser.add_argument(
        "--include-timestamps",
        type=parse_bool,
        default=True,
        help="Include timestamps in transcript mode (true/false, default: true)",
    )
    parser.add_argument(
        "--analysis-type",
        choices=["summary", "key_points", "quotes"],
        default="summary",
        help="Analysis style for analyze mode (default: summary)",
    )
    parser.add_argument(
        "--max-items",
        type=int,
        default=5,
        help="Maximum items in analysis output (default: 5)",
    )
    parser.add_argument(
        "--json",
        type=parse_bool,
        default=True,
        help="Emit JSON output (true/false, default: true)",
    )
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if args.max_items < 1:
        error = ToolError("invalid_argument", "--max-items must be >= 1", "Use a value like 3, 5, or 10.")
        emit(error_payload(error), as_json=args.json)
        return 1

    try:
        payload = run(args)
        emit(payload, as_json=args.json)
        return 0
    except ToolError as error:
        emit(error_payload(error), as_json=args.json)
        return 1
    except Exception as error:  # pragma: no cover - defensive fallback
        mapped = map_exception(error)
        emit(error_payload(mapped), as_json=args.json)
        return 1


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