import imageio
import os
import re
import tempfile
import yt_dlp

from datetime import timedelta
from google import genai
from google.genai import types

from smolagents import tool
from typing import List, Optional
from youtube_transcript_api import YouTubeTranscriptApi


# YouTube Video Review Tool
@tool
def review_youtube_video(url: str, question: str) -> str:
    """Reviews a YouTube video and answers a specific question about that video.
    Args:
        url (str): the URL to the YouTube video.
        question (str): The question you are asking about the video
    Returns:
        str: The answer to the question
    """
    try:
        client = genai.Client(api_key=os.getenv("GEMINI_KEY"))
        model = "models/gemini-1.5-flash-8b"

        response = client.models.generate_content(
            model=model,
            contents=types.Content(
                parts=[
                    types.Part(file_data=types.FileData(file_uri=url)),
                    types.Part(text=question),
                ]
            ),
        )
        return response.text
    except Exception as e:
        return f"Error asking {model} about video: {str(e)}"


@tool
def use_vision_model(
    question: str, image_paths: List[str], mime_type: str
) -> str:
    """Use a Vision Model to answer a question about a set of images.
    Args:
        question (str): The question you are asking about the images.
        image_paths (List[str]): The paths to the images to use for the question.
        mime_type (str): The mime type of the image.
    Returns:
        str: The answer to the question
    """
    try:
        client = genai.Client(api_key=os.getenv("GEMINI_KEY"))
        model = "models/gemini-2.0-flash-001"

        # Prepare the content parts
        parts = []
        for image_path in image_paths:
            with open(image_path, "rb") as f:
                image_bytes = f.read()

        response = []

        for chunk in client.models.generate_content_stream(
            model=model,
            contents=[
                question,
                types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
            ],
        ):
            response.append(chunk.text)

        return " ".join(response)

    except Exception as e:
        return f"Error using vision model: {str(e)}"


# YouTube Frames to Images Tool
@tool
def video_frames_to_images(
    url: str,
    folder_name: str,
    sample_interval_seconds: int = 5,
) -> List[str]:
    """Extracts frames from a video at specified intervals and saves them as images.
    Args:
        url (str): the URL to the video.
        folder_name (str): the name of the folder to save the images to.
        sample_interval_seconds (int): the interval between frames to sample.
    Returns:
        List[str]: A list of paths to the saved image files.
    """
    # Create a subdirectory for the frames
    frames_dir = os.path.join(folder_name, "frames")
    os.makedirs(frames_dir, exist_ok=True)

    ydl_opts = {
        "format": "bestvideo[height<=1080]+bestaudio/best[height<=1080]/best",
        "outtmpl": os.path.join(folder_name, "video.%(ext)s"),
        "quiet": True,
        "noplaylist": True,
        "merge_output_format": "mp4",
        "force_ipv4": True,
    }

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            video_path = next(
                (
                    os.path.join(folder_name, f)
                    for f in os.listdir(folder_name)
                    if f.endswith(".mp4")
                ),
                None,
            )

            if not video_path:
                raise RuntimeError("Failed to download video as mp4")

            reader = imageio.get_reader(video_path)
            metadata = reader.get_meta_data()
            fps = metadata.get("fps")

            if fps is None:
                reader.close()
                raise RuntimeError(
                    "Unable to determine FPS from video metadata"
                )

            frame_interval = int(fps * sample_interval_seconds)
            image_paths: List[str] = []

            for idx, frame in enumerate(reader):
                if idx % frame_interval == 0:
                    # Save frame as image
                    image_path = os.path.join(
                        frames_dir, f"frame_{idx:06d}.jpg"
                    )
                    imageio.imwrite(image_path, frame)
                    image_paths.append(image_path)

            reader.close()
            return image_paths

    except Exception as e:
        raise RuntimeError(f"Error processing video frames: {str(e)}") from e


@tool
def transcribe_youtube(url: str) -> str:
    """Transcribes a YouTube video using YouTube Transcript API or Gemini as fallback.
    Args:
        url (str): the URL to the YouTube video.
    Returns:
        str: The transcript of the YouTube video.
    """
    try:
        # First try using YouTube Transcript API
        video_id = _extract_video_id(url)
        if not video_id:
            raise ValueError(f"Invalid YouTube URL: {url}")

        try:
            # Try to get transcript in English
            transcript_chunks = YouTubeTranscriptApi.get_transcript(
                video_id, languages=["en"]
            )
            # Combine all chunks into a single transcript with timestamps
            transcript = ""
            for chunk in transcript_chunks:
                timestamp = str(timedelta(seconds=int(chunk["start"])))
                transcript += f"[{timestamp}] {chunk['text']}\n"
            return transcript

        except Exception as transcript_error:
            print(
                f"Failed to get transcript using YouTube API: {str(transcript_error)}"
            )
            print("Falling back to Gemini-based transcription...")

            # Fallback to Gemini-based transcription
            with tempfile.TemporaryDirectory() as tmpdir:
                # Download audio from YouTube
                ydl_opts = {
                    "format": "bestaudio/best",
                    "outtmpl": os.path.join(tmpdir, "audio.%(ext)s"),
                    "quiet": True,
                    "noplaylist": True,
                    "postprocessors": [
                        {
                            "key": "FFmpegExtractAudio",
                            "preferredcodec": "wav",
                            "preferredquality": "192",
                        }
                    ],
                }

                try:
                    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                        info = ydl.extract_info(url, download=True)
                        audio_path = next(
                            (
                                os.path.join(tmpdir, f)
                                for f in os.listdir(tmpdir)
                                if f.endswith(".wav")
                            ),
                            None,
                        )

                        if not audio_path:
                            raise RuntimeError(
                                "Failed to download audio"
                            ) from transcript_error

                        # Use Gemini to transcribe the audio
                        client = genai.Client(api_key=os.getenv("GEMINI_KEY"))
                        model = "models/gemini-1.5-flash-8b"

                        # Read the audio file
                        with open(audio_path, "rb") as audio_file:
                            audio_data = audio_file.read()

                        # Create the content with audio data
                        contents = types.Content(
                            parts=[
                                types.Part(
                                    file_data=types.FileData(
                                        mime_type="audio/wav",
                                        data=audio_data,
                                    )
                                ),
                                types.Part(
                                    text="Please transcribe this audio file. Include timestamps if possible."
                                ),
                            ]
                        )

                        # Generate transcription
                        response = client.models.generate_content(
                            model=model, contents=contents
                        )
                        return response.text

                except yt_dlp.utils.DownloadError as e:
                    raise RuntimeError(
                        f"Error downloading YouTube video: {str(e)}"
                    ) from transcript_error
                except Exception as e:
                    raise RuntimeError(
                        f"Error processing YouTube video: {str(e)}"
                    ) from transcript_error

    except Exception as e:
        raise RuntimeError(f"Error in YouTube transcription: {str(e)}") from e


def _extract_video_id(url: str) -> Optional[str]:
    """Extract video ID from YouTube URL.
    Args:
        url (str): the URL to the YouTube video.
    Returns:
        str: The video ID of the YouTube video.
    """
    patterns = [
        r"(?:youtube\.com\/watch\?v=|youtube\.com\/embed\/|youtu\.be\/)([^&\n?#]+)",
        r"(?:youtube\.com\/v\/|youtube\.com\/e\/|youtube\.com\/user\/[^\/]+\/|youtube\.com\/[^\/]+\/|youtube\.com\/embed\/|youtu\.be\/)([^&\n?#]+)",
    ]

    for pattern in patterns:
        match = re.search(pattern, url)
        if match:
            return match.group(1)
    return None