"""
Pydantic AI integration for Latch SDK

This module provides two integration patterns:

1. Context Manager (Recommended):
    from latch import latch
    from latch.integrations.pydantic_ai import get_model

    async with latch(api_key="...", task="legal-research") as ctx:
        model = get_model(ctx.model_name)
        agent = Agent(model=model, system_prompt="...")
        result = await agent.run("query")
        ctx.record_usage(input_tokens=result.usage.input_tokens,
                        output_tokens=result.usage.output_tokens)

2. Proxy Model (Legacy):
    from pydantic_ai import Agent
    from latch.integrations.pydantic_ai import LatchModel

    model = LatchModel(api_key="...", task="qa")
    agent = Agent(model=model, system_prompt="...")
    result = await agent.run("What is 2 + 2?")
"""

from __future__ import annotations

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, NotRequired, TypedDict

import httpx

from pydantic_ai.messages import (
    ModelMessage,
    ModelResponse,
    TextPart,
    ModelResponseStreamEvent,
    FinalResultEvent,
)
from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import RequestUsage
from pydantic_ai._run_context import RunContext


class LatchModelSettings(TypedDict, total=False):
    """Settings specific to Latch model (extends ModelSettings)"""

    task: str
    budget: float
    max_latency: int
    latch_metadata: dict[str, Any]
    max_tokens: int
    timeout: float | int


@dataclass
class LatchModel(Model):
    """
    Pydantic AI Model implementation that routes requests through Latch.

    Args:
        api_key: Your Latch API key
        base_url: Latch API base URL (default: https://api.latch.dev)
        task: Default task type for routing decisions
        budget: Default budget limit per request
        max_latency: Default maximum latency in milliseconds
        timeout: HTTP timeout in seconds
    """

    api_key: str
    base_url: str = "http://localhost:8081"
    task: str = "general"
    budget: float | None = None
    max_latency: int | None = None
    timeout: float = 60.0
    _client: httpx.AsyncClient = field(init=False, repr=False)

    def __post_init__(self) -> None:
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=self.timeout,
            headers={"Authorization": f"Bearer {self.api_key}"},
        )

    @property
    def model_name(self) -> str:
        return f"latch:{self.task}"

    @property
    def system(self) -> str:
        return "latch"

    def _extract_text_content(self, content: Any) -> str | None:
        """Extract text content from various content formats"""
        if isinstance(content, str):
            return content
        if isinstance(content, (list, tuple)):
            for item in content:
                if isinstance(item, str):
                    return item
                if isinstance(item, dict):
                    if item.get("type") == "text":
                        return item.get("text")
                if hasattr(item, "text"):
                    return item.text
        return None

    def _convert_messages(self, messages: list[ModelMessage]) -> list[dict[str, str]]:
        """Convert pydantic-ai messages to Latch format (OpenAI-compatible)"""
        result = []
        for msg in messages:
            if not hasattr(msg, "parts"):
                continue

            for part in msg.parts:
                part_kind = getattr(part, "part_kind", None)
                if part_kind is None:
                    continue

                raw_content = getattr(part, "content", None)
                if raw_content is None:
                    continue

                content = self._extract_text_content(raw_content)
                if content is None:
                    continue

                if part_kind == "system-prompt":
                    result.append({"role": "system", "content": content})
                elif part_kind == "user-prompt":
                    result.append({"role": "user", "content": content})
                elif part_kind == "text":
                    result.append({"role": "assistant", "content": content})

        return result

    def _get_settings(
        self, model_settings: ModelSettings | None
    ) -> tuple[str, float | None, int | None, dict | None]:
        """Extract Latch-specific settings from model_settings dict"""
        task = self.task
        budget = self.budget
        max_latency = self.max_latency
        metadata = None

        if model_settings and isinstance(model_settings, dict):
            if "task" in model_settings and model_settings["task"]:
                task = model_settings["task"]
            if "budget" in model_settings and model_settings["budget"] is not None:
                budget = model_settings["budget"]
            if "max_latency" in model_settings and model_settings["max_latency"] is not None:
                max_latency = model_settings["max_latency"]
            if "latch_metadata" in model_settings and model_settings["latch_metadata"] is not None:
                metadata = model_settings["latch_metadata"]

        return task, budget, max_latency, metadata

    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
        run_context: RunContext[Any] | None = None,
    ) -> ModelResponse:
        """Execute a request through Latch"""
        task, budget, max_latency, metadata = self._get_settings(model_settings)
        latch_messages = self._convert_messages(messages)

        if not latch_messages:
            part_kinds = []
            for msg in messages:
                if hasattr(msg, "parts"):
                    for part in msg.parts:
                        pk = getattr(part, "part_kind", None)
                        ct = type(getattr(part, "content", None)).__name__
                        part_kinds.append(f"{pk}:{ct}")
            raise ValueError(
                f"No convertible messages found. Input had {len(messages)} messages "
                f"with parts: {part_kinds}"
            )

        request_body = {
            "messages": latch_messages,
            "task": task,
        }
        if budget is not None:
            request_body["budget"] = budget
        if max_latency is not None:
            request_body["max_latency"] = max_latency
        if metadata is not None:
            request_body["metadata"] = metadata

        response = await self._client.post("/v1/execute", json=request_body)

        if response.status_code != 200:
            error_msg = f"Latch API error: {response.status_code}"
            try:
                error_data = response.json()
                if "error" in error_data:
                    error_msg = error_data["error"].get("message", error_msg)
            except Exception:
                pass
            raise RuntimeError(error_msg)

        data = response.json()["data"]

        usage = RequestUsage(
            input_tokens=data["usage"]["inputTokens"],
            output_tokens=data["usage"]["outputTokens"],
        )

        content = data["response"]["content"]

        return ModelResponse(
            parts=[TextPart(content=content)],
            usage=usage,
            model_name=data["model"],
            timestamp=datetime.now(timezone.utc),
        )

    @asynccontextmanager
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
        run_context: RunContext[Any] | None = None,
    ) -> AsyncIterator[LatchStreamedResponse]:
        """
        Stream a request through Latch.

        Note: Currently falls back to non-streaming since Latch execute
        endpoint doesn't support streaming yet. This provides API compatibility.
        """
        response = await self.request(messages, model_settings, model_request_parameters, run_context)
        yield LatchStreamedResponse(response)

    async def close(self) -> None:
        """Close the HTTP client"""
        await self._client.aclose()


class LatchStreamedResponse(StreamedResponse):
    """Simulated streaming response from Latch (wraps non-streaming response)"""

    def __init__(self, response: ModelResponse) -> None:
        self._response = response
        self._used = False

    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        """Yield a single final result event since Latch doesn't stream"""
        if not self._used:
            self._used = True
            content = ""
            for part in self._response.parts:
                if isinstance(part, TextPart):
                    content = part.content
                    break
            yield FinalResultEvent(content=content, tool_name=None)

    @property
    def provider_name(self) -> str:
        return "latch"

    @property
    def provider_url(self) -> str:
        return "https://latch.dev"

    def usage(self) -> RequestUsage:
        return self._response.usage

    def model_name(self) -> str:
        return self._response.model_name or "latch"

    def timestamp(self) -> datetime:
        return self._response.timestamp


def get_model(model_name: str) -> Model:
    """
    Create a pydantic-ai Model from a Latch model name.

    Use this with the context manager pattern:

        async with latch(task="research") as ctx:
            model = get_model(ctx.model_name)
            agent = Agent(model=model, ...)

    Args:
        model_name: Model name from Latch decision (e.g., "gpt-4o-mini", "claude-sonnet-4-5-20250514")

    Returns:
        A pydantic-ai Model instance

    Raises:
        ValueError: If the model name is not recognized
    """
    from pydantic_ai.models.openai import OpenAIModel
    from pydantic_ai.models.anthropic import AnthropicModel

    model_lower = model_name.lower()

    openai_prefixes = ("gpt", "o1", "o3", "o4", "davinci", "babbage")
    if any(model_lower.startswith(prefix) for prefix in openai_prefixes):
        return OpenAIModel(model_name)

    if "claude" in model_lower:
        return AnthropicModel(model_name)

    if "gemini" in model_lower:
        from pydantic_ai.models.gemini import GeminiModel
        return GeminiModel(model_name)

    raise ValueError(
        f"Unknown model: {model_name}. "
        "Supported prefixes: gpt, o1, o3, o4, claude, gemini"
    )
