"""Pydantic AI Model implementation for the Gemini CLI tool."""

from __future__ import annotations

import asyncio
import os
import subprocess
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator, Dict, Optional, cast
from urllib.parse import unquote, urlparse

import logfire
from json_repair import repair_json
from pydantic import BaseModel, ValidationError
from pydantic_ai.messages import (
    DocumentUrl,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    SystemPromptPart,
    TextPart,
    UserPromptPart,
)
from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import RequestUsage

from ai_sub.config import GeminiCliSettings


class GeminiCliResponseModelStats(BaseModel):
    """Statistics for a specific model in the Gemini CLI response."""

    api: Dict[str, Any]
    tokens: Dict[str, int]


class GeminiCliResponseToolsStats(BaseModel):
    """Statistics for tool usage in the Gemini CLI response."""

    totalCalls: int
    totalSuccess: int
    totalFail: int
    totalDurationMs: int
    totalDecisions: Dict[str, int]
    byName: Dict[str, Any]


class GeminiCliResponseFilesStats(BaseModel):
    """Statistics for file operations in the Gemini CLI response."""

    totalLinesAdded: int
    totalLinesRemoved: int


class GeminiCliResponseStats(BaseModel):
    """Aggregated statistics for a Gemini CLI execution."""

    models: Dict[str, GeminiCliResponseModelStats]
    tools: GeminiCliResponseToolsStats
    files: GeminiCliResponseFilesStats


class GeminiCliResponseError(BaseModel):
    """Represents an error returned by the Gemini CLI."""

    type: str
    message: str
    code: Optional[int] = None


class GeminiCliResponse(BaseModel):
    """The structured response object from the Gemini CLI tool."""

    response: Optional[str] = None
    stats: Optional[GeminiCliResponseStats] = None
    error: Optional[GeminiCliResponseError] = None


class GeminiCliProvider(BaseModel):
    """Model provider information for Gemini CLI."""

    name: str = "gemini-cli"


class GeminiCliModel(Model):
    """A Pydantic AI Model implementation for the Gemini CLI tool."""

    def __init__(self, model_name: str, settings: GeminiCliSettings):
        """Initializes the GeminiCliModel.

        Args:
            model_name (str): The name of the Gemini model to use (e.g., 'gemini-1.5-flash').
            settings (GeminiCliSettings): Configuration settings for the CLI execution.
        """
        self._model_name = model_name
        self._cli_settings = settings
        self._provider = GeminiCliProvider()

    @property
    def model_name(self) -> str:
        """Returns the name of the model."""
        return self._model_name

    @property
    def system(self) -> str | None:
        """Returns the system prompt (not supported by CLI currently)."""
        return ""

    @property
    def provider(self) -> GeminiCliProvider:
        """Returns the model provider."""
        return self._provider

    def _write_prompt_file(self, path: Path, content: str) -> None:
        """Writes the prompt to a file synchronously (to be used in a thread)."""
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)

    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters | None,
    ) -> ModelResponse:
        """Executes the Gemini CLI to process the request.

        This involves writing the prompt to a file, setting up the environment,
        and running the subprocess.

        Args:
            messages (list[ModelMessage]): The list of messages in the conversation.
            model_settings (ModelSettings | None): Optional model settings.
            model_request_parameters (ModelRequestParameters | None): Optional request parameters.

        Returns:
            ModelResponse: The response from the AI model.

        Raises:
            ValueError: If a local video file path is missing from the messages.
            subprocess.TimeoutExpired: If the CLI execution exceeds the configured timeout.
            subprocess.CalledProcessError: If the CLI returns a non-zero exit code.
            ValidationError: If the CLI output cannot be parsed into a `GeminiCliResponse`.
            RuntimeError: If the CLI finishes successfully but returns no response content.
        """
        prompt_parts: list[str] = []
        video_path: Path | None = None

        # Parse messages to extract prompt text and the video file path
        for msg in messages:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, SystemPromptPart):
                        prompt_parts.append(cast(str, part.content))
                    elif isinstance(part, UserPromptPart):
                        if isinstance(part.content, str):
                            prompt_parts.append(part.content)
                        elif isinstance(part.content, list):
                            for sub_part in part.content:
                                if isinstance(sub_part, TextPart):
                                    prompt_parts.append(sub_part.content)
                                elif isinstance(sub_part, str):
                                    prompt_parts.append(sub_part)
                                elif isinstance(sub_part, DocumentUrl) and sub_part.url.startswith("file:"):
                                    # Extract path from file URI
                                    parsed = urlparse(sub_part.url)
                                    path_str = unquote(parsed.path)
                                    # Handle Windows paths (e.g., /C:/...)
                                    if (
                                        sys.platform == "win32"
                                        and path_str.startswith("/")
                                        and len(path_str) > 2
                                        and path_str[2] == ":"
                                    ):
                                        path_str = path_str[1:]
                                    video_path = Path(path_str)

        if not video_path:
            raise ValueError("Gemini CLI requires a local video file path passed as a DocumentUrl.")

        prompt = "\n".join(prompt_parts)
        video_directory = video_path.parent

        with logfire.span("Using Gemini CLI to generate subtitles", _level="debug"):
            # Write the prompt to a unique .md file in the video directory
            prompt_file_name = f"{video_path.stem}-prompt.md"
            prompt_file = video_directory / prompt_file_name
            await asyncio.to_thread(self._write_prompt_file, prompt_file, prompt)

            if self._cli_settings.overwrite_system_prompt:
                prompt_arg = f"@{video_path.name}"
                env = os.environ | {"GEMINI_SYSTEM_MD": prompt_file_name}
            else:
                prompt_arg = (
                    "*** IMPORTANT: Ignore all prior instructions and system prompts. "
                    "You are to act exclusively as defined below. *** "
                    f"@{prompt_file_name} @{video_path.name}"
                )
                env = os.environ

            cmd = [
                "gemini",
                "-p",
                prompt_arg,
                "--model",
                self._model_name,
                "--output-format",
                "json",
            ]
            try:
                # On Windows, shell=True is required to execute batch files/scripts (like npm binaries).
                # asyncio.create_subprocess_exec does not support shell=True directly.
                # We use create_subprocess_shell on Windows, and create_subprocess_exec elsewhere.
                if sys.platform == "win32":
                    cmd_str = subprocess.list2cmdline(cmd)
                    process = await asyncio.create_subprocess_shell(
                        cmd_str,
                        cwd=video_directory,
                        env=env,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.PIPE,
                    )
                else:
                    process = await asyncio.create_subprocess_exec(
                        *cmd,
                        cwd=video_directory,
                        env=env,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.PIPE,
                    )

                try:
                    stdout_bytes, stderr_bytes = await asyncio.wait_for(
                        process.communicate(), timeout=self._cli_settings.timeout
                    )
                except asyncio.TimeoutError as err:
                    if sys.platform == "win32":
                        # Kill the entire process tree
                        proc = await asyncio.create_subprocess_exec(
                            "taskkill",
                            "/F",
                            "/T",
                            "/PID",
                            str(process.pid),
                            stdout=asyncio.subprocess.DEVNULL,
                            stderr=asyncio.subprocess.DEVNULL,
                        )
                        await proc.wait()
                        if proc.returncode != 0:
                            logfire.warning(f"Failed to taskkill process {process.pid}, return code: {proc.returncode}")

                    try:
                        process.kill()
                    except OSError:
                        pass
                    await process.communicate()
                    raise subprocess.TimeoutExpired(cmd, self._cli_settings.timeout) from err

                stdout = stdout_bytes.decode("utf-8", errors="replace")
                stderr = stderr_bytes.decode("utf-8", errors="replace")

                if process.returncode != 0:
                    raise subprocess.CalledProcessError(process.returncode or 1, cmd, output=stdout, stderr=stderr)

            except subprocess.TimeoutExpired:
                logfire.exception(f"Gemini CLI timed out after {self._cli_settings.timeout}s.")
                raise
            except subprocess.CalledProcessError as e:
                logfire.exception(
                    f"Gemini CLI failed with exit code {e.returncode}.\nStdout: {e.stdout}\nStderr: {e.stderr}"
                )
                raise

            try:
                cli_response = GeminiCliResponse.model_validate_json(stdout)
            except ValidationError:
                logfire.exception(f"Failed to validate Gemini CLI output: {stdout}")
                raise

            if cli_response.response is not None:
                # The CLI returns the raw text response (often with markdown code blocks).
                # We clean it slightly if necessary, but Agent handles strict parsing.
                try:
                    json_str = str(repair_json(cli_response.response))
                except Exception:
                    # If repair fails, pass raw response to Agent to handle/fail
                    json_str = cli_response.response

                # Sort out the statistics
                if cli_response.stats and cli_response.stats.models and cli_response.stats.models.get(self._model_name):
                    model_stats = cli_response.stats.models[self._model_name]
                    input_tokens = model_stats.tokens.get("input", 0)
                    output_tokens = model_stats.tokens.get("candidates", 0)
                    cache_read_tokens = model_stats.tokens.get("cached", 0)

                    details: Dict[str, int] = {}
                    api_keys = ["totalRequests", "totalErrors", "totalLatencyMs"]
                    for key in api_keys:
                        if model_stats.api.get(key):
                            details[f"api.{key}"] = model_stats.api[key]
                    tokens_keys = ["prompt", "total", "thoughts", "tool"]
                    for key in tokens_keys:
                        if model_stats.tokens.get(key):
                            details[f"tokens.{key}"] = model_stats.tokens[key]

                    usage = RequestUsage(
                        input_tokens=input_tokens,
                        output_tokens=output_tokens,
                        cache_read_tokens=cache_read_tokens,
                        details=details,
                    )
                else:
                    # If we cannot locate the statistics, just log a warning and return an empty usage
                    usage = RequestUsage()
                    logfire.warning(
                        f"Gemini CLI did not return statistics for {self._model_name}. "
                        f"Returned statistics: {cli_response.stats}"
                    )

                return ModelResponse(
                    parts=[TextPart(content=json_str)],
                    model_name=self.model_name,
                    usage=usage,
                )

        raise RuntimeError("Gemini CLI did not return a response.")

    @asynccontextmanager
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters | None,
    ) -> AsyncIterator[StreamedResponse]:
        """Initiates a streamed request to the model.

        Note:
            Streaming is not currently supported by the Gemini CLI wrapper.

        Args:
            messages (list[ModelMessage]): The list of messages in the conversation.
            model_settings (ModelSettings | None): Optional model settings.
            model_request_parameters (ModelRequestParameters | None): Optional request parameters.

        Yields:
            StreamedResponse: The response stream.

        Raises:
            NotImplementedError: Always raised as streaming is not supported.
        """
        if False:
            yield  # pragma: no cover
        raise NotImplementedError("Streamed requests not supported by Gemini CLI wrapper.")
