"""Provider backed by pydantic-ai for multi-model LLM support.

Supports model strings like ``openai:gpt-4o-mini``, ``anthropic:claude-sonnet``,
``gemini:gemini-2.0-flash``, etc.  Requires ``pydantic-ai`` to be installed::

    pip install pydantic-ai
"""

from __future__ import annotations

import asyncio
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any, ClassVar

try:
    from pydantic_ai import Agent

    _HAS_PYDANTIC_AI = True
except ImportError:  # pragma: no cover
    _HAS_PYDANTIC_AI = False

from parsantic.config import DEFAULT_MODEL

from .base import InferenceRequest
from .registry import register

# Match any "provider:model" string for common providers, plus bare model names.
_PATTERNS = (
    r"^(openai|anthropic|gemini|ollama|vertex|mistral|groq|bedrock|deepseek):",
    r"^gpt-",
    r"^claude-",
    r"^gemini-",
)


def _parse_model_spec(model_spec: str) -> tuple[str, str]:
    """Split ``'provider:model_name'`` into (provider, model_name).

    Bare model names (e.g. ``gpt-4o``) are mapped to their default provider.
    """
    if ":" in model_spec:
        provider, model_name = model_spec.split(":", 1)
        return provider, model_name
    if model_spec.startswith("gpt-"):
        return "openai", model_spec
    if model_spec.startswith("claude-"):
        return "anthropic", model_spec
    if model_spec.startswith("gemini-"):
        return "gemini", model_spec
    return "openai", model_spec


def _build_model_with_credentials(
    model_spec: str,
    api_key: str | None,
    base_url: str | None,
) -> Any:
    """Create a pydantic-ai Model object with explicit credentials.

    Falls back to the plain ``model_spec`` string when the provider-specific
    classes are unavailable (letting pydantic-ai resolve env vars).
    """
    provider_name, model_name = _parse_model_spec(model_spec)

    provider_kwargs: dict[str, Any] = {}
    if api_key:
        provider_kwargs["api_key"] = api_key
    if base_url:
        provider_kwargs["base_url"] = base_url

    def _build_provider(provider_cls: type[Any], kwargs: dict[str, Any]) -> Any:
        try:
            return provider_cls(**kwargs)
        except TypeError:
            # Some provider classes may not accept ``base_url``.
            if "base_url" in kwargs:
                reduced = {k: v for k, v in kwargs.items() if k != "base_url"}
                return provider_cls(**reduced)
            raise

    try:
        if provider_name == "openai":
            from pydantic_ai.models.openai import OpenAIModel
            from pydantic_ai.providers.openai import OpenAIProvider

            provider = _build_provider(OpenAIProvider, provider_kwargs)
            return OpenAIModel(model_name, provider=provider)

        if provider_name == "anthropic":
            from pydantic_ai.models.anthropic import AnthropicModel
            from pydantic_ai.providers.anthropic import AnthropicProvider

            provider = _build_provider(AnthropicProvider, provider_kwargs)
            return AnthropicModel(model_name, provider=provider)

        if provider_name == "gemini":
            from pydantic_ai.models.google import GoogleModel

            try:
                from pydantic_ai.providers.google import GoogleProvider
            except ImportError:
                from pydantic_ai.providers.google_gla import GoogleGLAProvider as GoogleProvider

            gla_kwargs = {k: v for k, v in provider_kwargs.items() if k == "api_key"}
            provider = _build_provider(GoogleProvider, gla_kwargs)
            return GoogleModel(model_name, provider=provider)
    except (ImportError, TypeError):
        pass

    # Unknown provider or import failure — let pydantic-ai resolve via env vars.
    return model_spec


@register(*_PATTERNS, priority=10)
@dataclass(slots=True)
class PydanticAIProvider:
    """Provider that delegates to pydantic-ai for model execution.

    The provider asks for raw text (``output_type=str``) so that
    parsantic's own parsing/coercion/alignment pipeline handles
    the structured output extraction downstream.
    """

    model_id: str | None = None
    api_key: str | None = None
    base_url: str | None = None
    max_concurrency: int = 8
    supported_attachment_kinds: ClassVar[frozenset[str]] = frozenset({"image", "pdf"})
    _agent: Any = field(default=None, repr=False, init=False)

    def __post_init__(self) -> None:
        if not _HAS_PYDANTIC_AI:
            raise ImportError(
                "pydantic-ai is required for this model. Install with: pip install pydantic-ai"
            )
        model_spec = self.model_id or DEFAULT_MODEL

        if self.api_key or self.base_url:
            model = _build_model_with_credentials(model_spec, self.api_key, self.base_url)
            self._agent = Agent(model, output_type=str)
        else:
            self._agent = Agent(model_spec, output_type=str)

    def infer(self, batch_prompts: Sequence[str], **kwargs: Any) -> Sequence[str]:
        results: list[str] = []
        for prompt in batch_prompts:
            result = self._agent.run_sync(prompt, **kwargs)
            results.append(result.output)
        return results

    def _build_message_parts(self, request: InferenceRequest) -> list[Any]:
        """Convert InferenceRequest to pydantic-ai message parts (text + BinaryContent)."""
        from pydantic_ai.messages import BinaryContent

        from parsantic.extract.media.attachments import AttachmentKind

        parts: list[Any] = []
        for attachment in request.attachments:
            if isinstance(attachment.source, bytes):
                data = attachment.source
            else:
                data = attachment.source.read_bytes()

            media_type = attachment.mime_type
            if media_type is None:
                if attachment.kind is AttachmentKind.PDF:
                    media_type = "application/pdf"
                elif attachment.kind is AttachmentKind.IMAGE:
                    media_type = "image/png"
                else:
                    media_type = "application/octet-stream"

            parts.append(BinaryContent(data=data, media_type=media_type))

        parts.append(request.prompt)
        return parts

    def infer_media(self, batch: Sequence[InferenceRequest], **kwargs: Any) -> Sequence[str]:
        results: list[str] = []
        for request in batch:
            parts = self._build_message_parts(request)
            result = self._agent.run_sync(parts, **kwargs)
            results.append(result.output)
        return results

    async def ainfer(self, batch_prompts: Sequence[str], **kwargs: Any) -> Sequence[str]:
        concurrency = kwargs.pop("max_concurrency", self.max_concurrency)
        semaphore = asyncio.Semaphore(max(1, int(concurrency)))

        async def _run_prompt(prompt: str) -> str:
            async with semaphore:
                result = await self._agent.run(prompt, **kwargs)
                return result.output

        return list(await asyncio.gather(*(_run_prompt(prompt) for prompt in batch_prompts)))

    async def ainfer_media(self, batch: Sequence[InferenceRequest], **kwargs: Any) -> Sequence[str]:
        concurrency = kwargs.pop("max_concurrency", self.max_concurrency)
        semaphore = asyncio.Semaphore(max(1, int(concurrency)))

        async def _run_request(request: InferenceRequest) -> str:
            async with semaphore:
                parts = self._build_message_parts(request)
                result = await self._agent.run(parts, **kwargs)
                return result.output

        return list(await asyncio.gather(*(_run_request(request) for request in batch)))
