#!/usr/bin/env python3
"""
YouTube MCP Server & CLI Tool for OpenClaw.

Integrates YouTube search, video metadata, and free transcript extraction
(subtitles) via YouTube Data API v3 and youtube-transcript-api.

Features:
- youtube_search: Search videos by keyword with title, channel, duration, and views
- youtube_get_transcript: Fetch German and English subtitles with timestamps (100% free, zero quota)
- youtube_video_details: Fetch full video metadata and description
"""

import os
import sys
import re
import json
import argparse
from typing import Dict, Any, List, Optional
from urllib.parse import urlparse, parse_qs

# Environment
YOUTUBE_API_KEY = os.environ.get("YOUTUBE_API_KEY", "").strip()


def extract_video_id(url_or_id: str) -> Optional[str]:
    """Extract 11-character YouTube video ID from various URL formats or raw ID."""
    if not url_or_id:
        return None
    url_or_id = url_or_id.strip()
    # If already an 11-char alphanumeric ID
    if re.match(r"^[a-zA-Z0-9_-]{11}$", url_or_id):
        return url_or_id
    
    # Common URL formats:
    # https://www.youtube.com/watch?v=VIDEO_ID
    # https://youtu.be/VIDEO_ID
    # https://www.youtube.com/shorts/VIDEO_ID
    # https://www.youtube.com/embed/VIDEO_ID
    parsed = urlparse(url_or_id)
    if "youtube.com" in parsed.netloc:
        if parsed.path == "/watch":
            qs = parse_qs(parsed.query)
            if "v" in qs and qs["v"]:
                return qs["v"][0]
        elif parsed.path.startswith("/shorts/"):
            parts = parsed.path.split("/")
            if len(parts) >= 3 and len(parts[2]) == 11:
                return parts[2]
        elif parsed.path.startswith("/embed/"):
            parts = parsed.path.split("/")
            if len(parts) >= 3 and len(parts[2]) == 11:
                return parts[2]
    elif "youtu.be" in parsed.netloc:
        path = parsed.path.lstrip("/")
        # Path might have query params or /
        video_id = path.split("?")[0].split("/")[0]
        if len(video_id) == 11:
            return video_id

    # Regex fallback
    match = re.search(r"[a-zA-Z0-9_-]{11}", url_or_id)
    if match:
        return match.group(0)
    return None


def parse_iso_duration(duration_str: str) -> str:
    """Convert ISO 8601 duration (PT1H2M30S or PT5M12S) to human readable (01:02:30 or 05:12)."""
    if not duration_str or not duration_str.startswith("PT"):
        return duration_str or ""
    hours = 0
    minutes = 0
    seconds = 0
    
    pattern = re.compile(r"PT(?:(\d+)H)?(?:(\d+)M)?(?:(\d+)S)?")
    match = pattern.match(duration_str)
    if match:
        h, m, s = match.groups()
        hours = int(h) if h else 0
        minutes = int(m) if m else 0
        seconds = int(s) if s else 0
        if hours > 0:
            return f"{hours}:{minutes:02d}:{seconds:02d}"
        return f"{minutes:02d}:{seconds:02d}"
    return duration_str


def format_seconds(seconds: float) -> str:
    """Format seconds into HH:MM:SS or MM:SS."""
    s = int(seconds)
    hours = s // 3600
    minutes = (s % 3600) // 60
    secs = s % 60
    if hours > 0:
        return f"{hours:02d}:{minutes:02d}:{secs:02d}"
    return f"{minutes:02d}:{secs:02d}"


# ==============================================================================
# YouTube Operations
# ==============================================================================

def search_videos(
    query: str,
    max_results: int = 5,
    video_duration: str = "any",
    api_key: Optional[str] = None,
) -> Dict[str, Any]:
    """Search YouTube videos using YouTube Data API v3."""
    key = api_key or YOUTUBE_API_KEY
    if not key:
        return {
            "success": False,
            "error": "YOUTUBE_API_KEY is not set. Please provide a valid Google Cloud API key.",
            "videos": [],
        }

    try:
        from googleapiclient.discovery import build
        youtube = build("youtube", "v3", developerKey=key, cache_discovery=False)
        
        search_kwargs = {
            "q": query,
            "part": "snippet",
            "type": "video",
            "maxResults": min(max(1, max_results), 20),
        }
        if video_duration in ["short", "medium", "long"]:
            search_kwargs["videoDuration"] = video_duration

        search_resp = youtube.search().list(**search_kwargs).execute()
        items = search_resp.get("items", [])
        if not items:
            return {"success": True, "query": query, "count": 0, "videos": []}

        video_ids = [it["id"]["videoId"] for it in items if "videoId" in it.get("id", {})]
        
        # Batch fetch details (duration, view count, stats) - only 1 unit!
        details_map = {}
        if video_ids:
            videos_resp = youtube.videos().list(
                id=",".join(video_ids),
                part="contentDetails,statistics,snippet"
            ).execute()
            for v in videos_resp.get("items", []):
                details_map[v["id"]] = v

        videos = []
        for item in items:
            vid = item["id"].get("videoId")
            if not vid:
                continue
            snip = item.get("snippet", {})
            v_detail = details_map.get(vid, {})
            cd = v_detail.get("contentDetails", {})
            stats = v_detail.get("statistics", {})

            duration_raw = cd.get("duration", "")
            duration_formatted = parse_iso_duration(duration_raw)

            videos.append({
                "video_id": vid,
                "title": snip.get("title", ""),
                "channel": snip.get("channelTitle", ""),
                "published_at": snip.get("publishedAt", ""),
                "duration": duration_formatted,
                "view_count": stats.get("viewCount", "N/A"),
                "url": f"https://www.youtube.com/watch?v={vid}",
                "description_snippet": snip.get("description", "")[:200],
                "thumbnail": snip.get("thumbnails", {}).get("high", {}).get("url", ""),
            })

        return {
            "success": True,
            "query": query,
            "count": len(videos),
            "videos": videos,
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"YouTube search failed: {str(e)}",
            "videos": [],
        }


def get_transcript(
    video_url_or_id: str,
    languages: Optional[List[str]] = None,
    include_timestamps: bool = True,
) -> Dict[str, Any]:
    """Fetch video transcript/subtitles using youtube-transcript-api."""
    video_id = extract_video_id(video_url_or_id)
    if not video_id:
        return {
            "success": False,
            "error": f"Could not extract a valid YouTube video ID from '{video_url_or_id}'.",
        }

    lang_pref = languages or ["de", "en"]

    try:
        from youtube_transcript_api import YouTubeTranscriptApi
        
        api = YouTubeTranscriptApi()
        
        # Try fetching with preferred languages
        transcript_obj = None
        detected_lang = None
        is_generated = False

        # Use .list(video_id) to inspect available tracks and find best match
        try:
            transcript_list = api.list(video_id)
            # Find best match or auto-generated
            try:
                transcript_obj = transcript_list.find_transcript(lang_pref)
            except Exception:
                # If no direct match in lang_pref, try any available transcript and translate if possible
                for t in transcript_list:
                    transcript_obj = t
                    break

            if transcript_obj:
                detected_lang = transcript_obj.language_code
                is_generated = transcript_obj.is_generated
                snippets = transcript_obj.fetch()
            else:
                # Fallback to direct fetch
                snippets = api.fetch(video_id, languages=lang_pref)
        except Exception:
            # Fallback to direct fetch
            snippets = api.fetch(video_id, languages=lang_pref)

        formatted_lines = []
        full_text_chunks = []
        total_duration = 0.0

        for item in snippets:
            # item may be an object (FetchedTranscriptSnippet) or dict
            if hasattr(item, "text"):
                text = item.text.strip()
                start = getattr(item, "start", 0.0)
                duration = getattr(item, "duration", 0.0)
            elif isinstance(item, dict):
                text = item.get("text", "").strip()
                start = item.get("start", 0.0)
                duration = item.get("duration", 0.0)
            else:
                text = str(item).strip()
                start = 0.0
                duration = 0.0

            if not text:
                continue

            timestamp = format_seconds(start)
            formatted_lines.append(f"[{timestamp}] {text}")
            full_text_chunks.append(text)
            total_duration = max(total_duration, start + duration)

        full_text = " ".join(full_text_chunks)

        return {
            "success": True,
            "video_id": video_id,
            "url": f"https://www.youtube.com/watch?v={video_id}",
            "language": detected_lang or "auto",
            "is_generated": is_generated,
            "duration_seconds": round(total_duration, 1),
            "duration_formatted": format_seconds(total_duration),
            "snippet_count": len(formatted_lines),
            "formatted_transcript": "\n".join(formatted_lines) if include_timestamps else full_text,
            "full_text": full_text,
        }
    except Exception as e:
        err_msg = str(e)
        if "TranscriptsDisabled" in err_msg or "Subtitles are disabled" in err_msg:
            return {
                "success": False,
                "video_id": video_id,
                "error": "Transcripts are disabled for this video by the uploader.",
            }
        elif "NoTranscriptFound" in err_msg:
            return {
                "success": False,
                "video_id": video_id,
                "error": f"No transcript found in requested languages ({lang_pref}).",
            }
        return {
            "success": False,
            "video_id": video_id,
            "error": f"Failed to retrieve transcript: {err_msg}",
        }


def get_video_details(
    video_url_or_id: str,
    api_key: Optional[str] = None,
) -> Dict[str, Any]:
    """Fetch complete video details and description."""
    video_id = extract_video_id(video_url_or_id)
    if not video_id:
        return {
            "success": False,
            "error": f"Could not extract a valid YouTube video ID from '{video_url_or_id}'.",
        }

    key = api_key or YOUTUBE_API_KEY
    if not key:
        return {
            "success": False,
            "error": "YOUTUBE_API_KEY is not set.",
        }

    try:
        from googleapiclient.discovery import build
        youtube = build("youtube", "v3", developerKey=key, cache_discovery=False)
        resp = youtube.videos().list(
            id=video_id,
            part="snippet,contentDetails,statistics"
        ).execute()

        items = resp.get("items", [])
        if not items:
            return {
                "success": False,
                "error": f"Video '{video_id}' not found.",
            }

        item = items[0]
        snip = item.get("snippet", {})
        cd = item.get("contentDetails", {})
        stats = item.get("statistics", {})

        return {
            "success": True,
            "video_id": video_id,
            "title": snip.get("title", ""),
            "channel": snip.get("channelTitle", ""),
            "published_at": snip.get("publishedAt", ""),
            "description": snip.get("description", ""),
            "duration": parse_iso_duration(cd.get("duration", "")),
            "view_count": stats.get("viewCount", "0"),
            "like_count": stats.get("likeCount", "0"),
            "comment_count": stats.get("commentCount", "0"),
            "tags": snip.get("tags", []),
            "url": f"https://www.youtube.com/watch?v={video_id}",
            "thumbnail": snip.get("thumbnails", {}).get("maxres", snip.get("thumbnails", {}).get("high", {})).get("url", ""),
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Failed to get video details: {str(e)}",
        }


# ==============================================================================
# MCP Protocol Handlers (JSON-RPC 2.0 stdio)
# ==============================================================================

TOOLS_DEFINITIONS = [
    {
        "name": "youtube_search",
        "description": "Search YouTube for videos matching a query. Returns video title, channel, video ID, formatted duration, view count, and direct link.",
        "inputSchema": {
            "type": "object",
            "required": ["query"],
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search term or topic (e.g. 'Ninja Creami Deluxe Rezepte', 'Kubernetes Helm Tutorial').",
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of videos to return (default 5, max 20).",
                    "default": 5,
                },
                "video_duration": {
                    "type": "string",
                    "description": "Filter by duration: 'any' (default), 'short' (< 4m), 'medium' (4m-20m), 'long' (> 20m).",
                    "enum": ["any", "short", "medium", "long"],
                    "default": "any",
                },
            },
        },
    },
    {
        "name": "youtube_get_transcript",
        "description": "Fetch video transcript / subtitles with timestamps. Free and unlimited. Essential for summarizing YouTube videos, extracting recipes, interview takeaways, or step-by-step instructions.",
        "inputSchema": {
            "type": "object",
            "required": ["video_url_or_id"],
            "properties": {
                "video_url_or_id": {
                    "type": "string",
                    "description": "YouTube URL (e.g. 'https://youtu.be/...' or 'https://www.youtube.com/watch?v=...') or 11-char video ID.",
                },
                "languages": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Preferred language codes in priority order (default ['de', 'en']).",
                    "default": ["de", "en"],
                },
                "include_timestamps": {
                    "type": "boolean",
                    "description": "Whether to include [HH:MM:SS] timestamps for each segment (default True).",
                    "default": True,
                },
            },
        },
    },
    {
        "name": "youtube_video_details",
        "description": "Fetch complete YouTube video details including full description, channel, publish date, duration, tags, and stats.",
        "inputSchema": {
            "type": "object",
            "required": ["video_url_or_id"],
            "properties": {
                "video_url_or_id": {
                    "type": "string",
                    "description": "YouTube video URL or 11-character video ID.",
                },
            },
        },
    },
]


def run_mcp_server():
    """Run JSON-RPC 2.0 stdio loop for OpenClaw."""
    while True:
        try:
            line = sys.stdin.readline()
            if not line:
                break
            line = line.strip()
            if not line:
                continue

            req = json.loads(line)
            req_id = req.get("id")
            method = req.get("method")
            params = req.get("params", {})

            if method == "initialize":
                res = {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {"tools": {}},
                        "serverInfo": {
                            "name": "openclaw-youtube",
                            "version": "1.0.0",
                        },
                    },
                }
                print(json.dumps(res), flush=True)

            elif method == "notifications/initialized":
                pass

            elif method == "tools/list":
                res = {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {"tools": TOOLS_DEFINITIONS},
                }
                print(json.dumps(res), flush=True)

            elif method == "tools/call":
                tool_name = params.get("name")
                args = params.get("arguments", {})
                content = []

                if tool_name == "youtube_search":
                    result = search_videos(
                        query=args.get("query", ""),
                        max_results=args.get("max_results", 5),
                        video_duration=args.get("video_duration", "any"),
                    )
                    content.append({"type": "text", "text": json.dumps(result, indent=2, ensure_ascii=False)})

                elif tool_name == "youtube_get_transcript":
                    result = get_transcript(
                        video_url_or_id=args.get("video_url_or_id", ""),
                        languages=args.get("languages", ["de", "en"]),
                        include_timestamps=args.get("include_timestamps", True),
                    )
                    content.append({"type": "text", "text": json.dumps(result, indent=2, ensure_ascii=False)})

                elif tool_name == "youtube_video_details":
                    result = get_video_details(
                        video_url_or_id=args.get("video_url_or_id", "")
                    )
                    content.append({"type": "text", "text": json.dumps(result, indent=2, ensure_ascii=False)})

                else:
                    content.append({"type": "text", "text": f"Error: Unknown tool '{tool_name}'"})

                res = {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {"content": content},
                }
                print(json.dumps(res), flush=True)

            else:
                res = {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "error": {"code": -32601, "message": f"Method '{method}' not found"},
                }
                print(json.dumps(res), flush=True)

        except Exception as e:
            err_res = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {"code": -32603, "message": f"Internal error: {str(e)}"},
            }
            print(json.dumps(err_res), flush=True)


def run_cli():
    """CLI testing interface."""
    parser = argparse.ArgumentParser(description="YouTube MCP Server & CLI Tool")
    subparsers = parser.add_subparsers(dest="command", help="Subcommand")

    # Search
    p_search = subparsers.add_parser("search", help="Search videos")
    p_search.add_argument("query", help="Search query")
    p_search.add_argument("--limit", "-n", type=int, default=5, help="Max results")
    p_search.add_argument("--duration", choices=["any", "short", "medium", "long"], default="any")

    # Transcript
    p_trans = subparsers.add_parser("transcript", help="Get transcript")
    p_trans.add_argument("video", help="URL or video ID")
    p_trans.add_argument("--languages", "-l", default="de,en", help="Comma separated language preferences")
    p_trans.add_argument("--no-timestamps", action="store_true", help="Omit timestamps")

    # Details
    p_details = subparsers.add_parser("details", help="Get video details")
    p_details.add_argument("video", help="URL or video ID")

    args = parser.parse_args()

    if args.command == "search":
        res = search_videos(query=args.query, max_results=args.limit, video_duration=args.duration)
        print(json.dumps(res, indent=2, ensure_ascii=False))
    elif args.command == "transcript":
        langs = [l.strip() for l in args.languages.split(",") if l.strip()]
        res = get_transcript(video_url_or_id=args.video, languages=langs, include_timestamps=not args.no_timestamps)
        print(json.dumps(res, indent=2, ensure_ascii=False))
    elif args.command == "details":
        res = get_video_details(video_url_or_id=args.video)
        print(json.dumps(res, indent=2, ensure_ascii=False))
    else:
        parser.print_help()


if __name__ == "__main__":
    if len(sys.argv) == 1 or (len(sys.argv) == 2 and sys.argv[1] == "--stdio"):
        run_mcp_server()
    else:
        run_cli()