#!/usr/bin/env python3

from __future__ import annotations

import re
import sys
import json
import logging
import time

from pathlib import Path
from typing import Optional, List
from urllib.parse import urlparse, parse_qs
# from google import genai

from youtube_transcript_api import (
    YouTubeTranscriptApi,
    TranscriptsDisabled,
    NoTranscriptFound,
    VideoUnavailable,
)


# ============================================================================
# CONFIG
# ============================================================================

BASE_DIR        = Path(__file__).resolve().parent
OUTPUT_DIR      = BASE_DIR / "output"
OUTPUT_DIR.mkdir(exist_ok=True)

TRANSCRIPT_FILE = OUTPUT_DIR / "transcript.txt"
SUMMARY_FILE    = OUTPUT_DIR / "summary.txt"
QA_FILE         = OUTPUT_DIR / "qa.txt"

GEMINI_API_KEY  = "AIzaSyCNz5wQAyJ65kNRkwr0-1A-_Z6-lQzdcyc"

GEMINI_MODELS = [
    "gemini-2.5-flash",
    "gemini-2.5-flash-lite",
    "gemini-2.5-pro",
]

POLLING_CONFIG = {
    "attempt_1": {"wait_before": 0,    "description": "Immediate attempt on trigger"},
    "attempt_2": {"wait_before": 300,  "description": "Retry after 5 minutes"},
    "attempt_3": {"wait_before": 900,  "description": "Retry after 15 minutes"},
    "attempt_4": {"wait_before": 1800, "description": "Retry after 30 minutes"},
    "attempt_5": {"wait_before": 3600, "description": "Retry after 1 hour"},
    "attempt_6": {"wait_before": 3600, "description": "Retry after 2 hours total"},
    "attempt_7": {"wait_before": 3600, "description": "Retry after 3 hours total"},
    "attempt_8": {"wait_before": 3600, "description": "Retry after 4 hours total"},
    "attempt_9": {"wait_before": 3600, "description": "Final attempt at 5 hours total"},
}

SYSTEM_PROMPT = """
You are an expert content summarizer and educator.

Produce the full output containing exactly two parts separated by a line with only 5 exclamation marks:

!!!!!

--- PART 1: SUMMARY ---

Write a detailed, well-structured summary of the entire content.
Use the following structure:

## Overview
A 3-5 sentence high-level overview of the entire content.

## Key Topics Covered
List the main topics discussed, each with a brief explanation.

## Detailed Summary
A thorough section-by-section breakdown of the content in the order it was presented.
Use subheadings for each major section or topic shift.

## Key Takeaways
A bullet list of the most important insights, facts, or conclusions from the content.

---

!!!!!

--- PART 2: Q&A ---

Generate a comprehensive Q&A section based on the content.
Format each entry exactly like this:

Q1: [First question]
Answer: [Detailed answer]

Q2: [Second question]
Answer: [Detailed answer]

Q3: [Third question]
Answer: [Detailed answer]

... and so on until all important questions are covered.

Rules:
- Number every question and answer with matching numbers (Q1/A1, Q2/A2, etc.)
- Each answer must be detailed and self-contained
- Cover all major topics, concepts, facts, and insights from the content
- Minimum 10 Q&A pairs, more if the content is rich
- Do NOT use bullet points inside answers — write in full sentences

---
"""


# ============================================================================
# LOGGING
# ============================================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
logger = logging.getLogger("gemini_pipeline")


# ============================================================================
# HELPERS
# ============================================================================

def _format_duration(seconds: int) -> str:
    if seconds < 60:
        return f"{seconds}s"
    if seconds < 3600:
        return f"{seconds // 60}m"
    h = seconds // 3600
    m = (seconds % 3600) // 60
    return f"{h}h {m}m" if m else f"{h}h"


def fetch_video_title(video_id: str) -> str:
    """Fetch YouTube video title via oembed — no API key needed."""
    try:
        import urllib.request
        url = (
            f"https://www.youtube.com/oembed"
            f"?url=https://www.youtube.com/watch?v={video_id}&format=json"
        )
        with urllib.request.urlopen(url, timeout=10) as resp:
            data  = json.loads(resp.read().decode())
            title = data.get("title", "")
            safe  = re.sub(r'[\\/*?:"<>|]', "", title)
            safe  = re.sub(r"\s+", "_", safe.strip())
            return safe[:80] or video_id
    except Exception:
        return video_id


# ============================================================================
# YOUTUBE TRANSCRIPT FETCHER
# ============================================================================

class YouTubeTranscriptFetcher:
    """Fetches YouTube transcript with polling retry for new uploads."""

    def __init__(
        self,
        youtube_url: str,
        output_file: Path              = TRANSCRIPT_FILE,
        languages: Optional[List[str]] = None,
        polling_config: dict           = None,
    ):
        self.youtube_url    = youtube_url
        self.output_file    = Path(output_file)
        self.languages      = languages or ["en", "en-US", "en-GB"]
        self.polling_config = polling_config or POLLING_CONFIG
        self.video_id       = self._extract_video_id(youtube_url)
        self.api            = YouTubeTranscriptApi()

    @staticmethod
    def _extract_video_id(url: str) -> str:
        parsed = urlparse(url)
        if parsed.hostname == "youtu.be":
            return parsed.path.lstrip("/").split("?")[0]
        if parsed.hostname in ("youtube.com", "www.youtube.com", "m.youtube.com"):
            path_parts = parsed.path.strip("/").split("/")
            if path_parts[0] in ("live", "shorts", "embed") and len(path_parts) >= 2:
                return path_parts[1].split("?")[0]
            params = parse_qs(parsed.query)
            if "v" in params:
                return params["v"][0]
            raise ValueError(f"Could not extract video ID from URL: {url}")
        raise ValueError(f"Unsupported YouTube URL: {url}")

    def _fetch_once(self) -> str:
        transcript = self.api.fetch(self.video_id, languages=self.languages)
        return " ".join(item.text for item in transcript)

    def _save(self, text: str) -> None:
        self.output_file.parent.mkdir(parents=True, exist_ok=True)
        self.output_file.write_text(text, encoding="utf-8")

    def run(self) -> str:
        logger.info("Video ID    : %s", self.video_id)
        logger.info("Output file : %s", self.output_file)
        logger.info("Total polling attempts: %d", len(self.polling_config))

        attempts = list(self.polling_config.items())

        for idx, (attempt_key, config) in enumerate(attempts, start=1):
            wait_before = config["wait_before"]
            description = config["description"]

            if wait_before > 0:
                logger.info(
                    "[%d/%d] %s — waiting %s before retry...",
                    idx, len(attempts), description,
                    _format_duration(wait_before),
                )
                time.sleep(wait_before)

            logger.info(
                "[%d/%d] %s — fetching transcript now...",
                idx, len(attempts), description,
            )

            try:
                text = self._fetch_once()
                self._save(text)
                logger.info(
                    "[%d/%d] ✅ Transcript fetched — %d characters",
                    idx, len(attempts), len(text),
                )
                return text

            except TranscriptsDisabled as e:
                logger.warning("[%d/%d] Transcripts disabled: %s", idx, len(attempts), e)
                raise  # no point retrying

            except VideoUnavailable as e:
                logger.warning("[%d/%d] Video unavailable: %s", idx, len(attempts), e)

            except NoTranscriptFound as e:
                logger.warning("[%d/%d] No transcript yet: %s", idx, len(attempts), e)

            except KeyboardInterrupt:
                logger.warning("Interrupted by user.")
                raise

            except Exception as e:
                logger.exception("[%d/%d] Unexpected error: %s", idx, len(attempts), e)

            if idx < len(attempts):
                next_cfg = attempts[idx][1]
                logger.info(
                    "[%d/%d] Will retry in %s (%s)",
                    idx, len(attempts),
                    _format_duration(next_cfg["wait_before"]),
                    next_cfg["description"],
                )
            else:
                logger.error("All %d polling attempts exhausted.", len(attempts))

        raise RuntimeError(
            f"Transcript not available after {len(attempts)} attempts (~5 hours). "
            f"Video ID: {self.video_id}"
        )


# ============================================================================
# GEMINI SUMMARIZER
# ============================================================================

class GeminiSummarizer:
    """Sends transcript to Gemini with model fallback + per-model retry."""

    # Retry config
    MAX_RETRIES  = 5
    BASE_WAIT    = 10    # seconds
    MAX_WAIT     = 120   # seconds cap

    # Errors → retry same model with backoff
    RETRYABLE    = ["503", "502", "500", "UNAVAILABLE", "SERVICE_UNAVAILABLE"]
    # Errors → skip to next model immediately
    SKIP_TO_NEXT = ["429", "RESOURCE_EXHAUSTED", "quota", "404", "NOT_FOUND"]

    def __init__(
        self,
        api_key: str       = GEMINI_API_KEY,
        models: list       = None,
        summary_file: Path = SUMMARY_FILE,
        qa_file: Path      = QA_FILE,
    ):
        self.client       = genai.Client(api_key=api_key)
        self.models       = models or GEMINI_MODELS
        self.summary_file = Path(summary_file)
        self.qa_file      = Path(qa_file)

    def _call_api(self, transcript: str) -> tuple[str, str]:
        """
        Try each model in order.
        Per model: retry up to MAX_RETRIES on transient errors with backoff.
        Returns (response_text, model_used).
        """
        overall_last_error = None

        for model in self.models:
            logger.info("── Trying model: %s", model)
            wait     = self.BASE_WAIT
            last_err = None

            for attempt in range(1, self.MAX_RETRIES + 1):
                try:
                    logger.info("   [%d/%d] Sending request...", attempt, self.MAX_RETRIES)
                    response = self.client.models.generate_content(
                        model=model,
                        contents=transcript,
                        config={"system_instruction": SYSTEM_PROMPT},
                    )
                    logger.info(
                        "✅ Response received from: %s (attempt %d)",
                        model, attempt,
                    )
                    return response.text, model

                except Exception as e:
                    err      = str(e)
                    last_err = e

                    if any(k in err for k in self.SKIP_TO_NEXT):
                        logger.warning(
                            "   [%d/%d] %s — quota/not-found, skipping to next model.",
                            attempt, self.MAX_RETRIES, model,
                        )
                        break  # skip to next model

                    elif any(k in err for k in self.RETRYABLE):
                        if attempt < self.MAX_RETRIES:
                            logger.warning(
                                "   [%d/%d] %s — transient error. "
                                "Retrying in %ds...",
                                attempt, self.MAX_RETRIES, model, wait,
                            )
                            time.sleep(wait)
                            wait = min(wait * 2, self.MAX_WAIT)
                        else:
                            logger.warning(
                                "   [%d/%d] %s — max retries reached, "
                                "trying next model.",
                                attempt, self.MAX_RETRIES, model,
                            )

                    else:
                        logger.error(
                            "   [%d/%d] %s — unhandled error: %s",
                            attempt, self.MAX_RETRIES, model, err,
                        )
                        raise

            overall_last_error = last_err

        raise RuntimeError(
            f"All models and retries exhausted. Last error: {overall_last_error}"
        )

    @staticmethod
    def _split(full_text: str) -> tuple[str, str]:
        for pattern in (r"^\s*!{5}\s*$", r"^\s*!{3}\s*$"):
            parts = re.split(pattern, full_text, flags=re.MULTILINE)
            if len(parts) >= 2:
                return parts[0].strip(), "".join(parts[1:]).strip()
        return full_text.strip(), ""

    def run(self, transcript: str) -> tuple[str, str, str]:
        full, model_used = self._call_api(transcript)
        summary, qa      = self._split(full)

        self.summary_file.write_text(summary, encoding="utf-8")
        self.qa_file.write_text(qa,           encoding="utf-8")

        logger.info("Summary saved → %s", self.summary_file)
        logger.info("Q&A saved    → %s", self.qa_file)

        return summary, qa, model_used


# ============================================================================
# PIPELINE
# ============================================================================

class TranscriptSummaryPipeline:

    def __init__(
        self,
        youtube_url: str,
        languages: Optional[List[str]] = None,
        polling_config: dict           = None,
    ):
        self.youtube_url = youtube_url
        self.fetcher     = YouTubeTranscriptFetcher(
            youtube_url=youtube_url,
            output_file=TRANSCRIPT_FILE,
            languages=languages,
            polling_config=polling_config,
        )
        self.summarizer  = GeminiSummarizer()
        self.video_id    = self.fetcher.video_id
        self.video_title = fetch_video_title(self.video_id)

    def run(self) -> dict:
        logger.info("=== Pipeline started ===")
        logger.info("Video title : %s", self.video_title)

        transcript         = self.fetcher.run()
        summary, qa, model = self.summarizer.run(transcript)

        logger.info("=== Pipeline complete | model: %s ===", model)

        return {
            "video_id":    self.video_id,
            "video_title": self.video_title,
            "model_used":  model,
            "summary":     summary,
            "qa":          qa,
            "transcript":  transcript,
        }


# ============================================================================
# CLI
# ============================================================================

def main():
    if len(sys.argv) < 2:
        print("Usage: python gemini.py <youtube_url>", file=sys.stderr)
        sys.exit(1)

    pipeline = TranscriptSummaryPipeline(
        youtube_url=sys.argv[1],
        languages=["en", "en-US", "en-GB"],
    )
    result = pipeline.run()
    for key, value in result.items():
        if key not in ("summary", "qa", "transcript"):
            print(f"{key}: {value}")


if __name__ == "__main__":
    main()