"""
PydanticAI provider for Ollama with message sanitization.

Ollama's OpenAI-compatible API rejects messages with `content: null`,
which PydanticAI sends for assistant messages that only contain tool calls.
This provider sanitizes messages to use empty strings instead of null.
"""
from typing import Any

from openai import AsyncOpenAI
from pydantic_ai.providers.ollama import OllamaProvider

from src.core.config import config
from src.core.logging_config import get_logger

logger = get_logger(__name__)


class TatlockOllamaProvider(OllamaProvider):
    """
    Custom OllamaProvider with message sanitization for Tatlock agents.

    Fixes the 'invalid message content type: <nil>' error that occurs
    when assistant messages have `content: null` with tool calls.
    """

    def __init__(self, base_url: str | None = None):
        """
        Initialize provider with Ollama base URL.

        Args:
            base_url: Ollama API URL (defaults to config.OLLAMA_HOST/v1)
        """
        if base_url is None:
            clean_host = str(config.OLLAMA_HOST).rstrip("/")
            base_url = f"{clean_host}/v1"

        super().__init__(base_url=base_url)

        # Override the client with our sanitized version
        self._openai_client = _SanitizedAsyncOpenAI(base_url=base_url)

        logger.debug("tatlock_ollama_provider_created", base_url=base_url)


class _SanitizedAsyncOpenAI(AsyncOpenAI):
    """AsyncOpenAI client that sanitizes messages before sending."""

    def __init__(self, **kwargs: Any):
        # Ollama doesn't need an API key
        super().__init__(api_key="ollama", **kwargs)

    @property
    def chat(self) -> "_SanitizedChat":
        """Return sanitized chat interface."""
        return _SanitizedChat(self)


class _SanitizedChat:
    """Chat interface wrapper with sanitized completions."""

    def __init__(self, client: _SanitizedAsyncOpenAI):
        self._client = client
        self._original_chat = AsyncOpenAI.chat.fget(client)  # type: ignore

    @property
    def completions(self) -> "_SanitizedCompletions":
        """Return sanitized completions interface."""
        return _SanitizedCompletions(self._original_chat.completions)


class _SanitizedCompletions:
    """Completions wrapper that sanitizes messages before API calls."""

    def __init__(self, original_completions: Any):
        self._original = original_completions

    async def create(self, **kwargs: Any) -> Any:
        """
        Create chat completion with sanitized messages.

        Converts `content: null` to `content: ""` in assistant messages
        to prevent Ollama's 'invalid message content type: <nil>' error.
        """
        if "messages" in kwargs:
            kwargs["messages"] = _sanitize_messages(kwargs["messages"])

        return await self._original.create(**kwargs)


def _sanitize_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """
    Sanitize messages to fix null content issues.

    When an assistant message has tool_calls but no text content,
    PydanticAI sets content to None. Ollama rejects this.
    We convert None to empty string.

    Args:
        messages: List of chat messages

    Returns:
        Sanitized messages with null content replaced by empty strings
    """
    sanitized = []
    for msg in messages:
        msg_copy = dict(msg)

        # Fix null content in assistant messages with tool calls
        if msg_copy.get("role") == "assistant":
            if msg_copy.get("content") is None and msg_copy.get("tool_calls"):
                msg_copy["content"] = ""
                logger.debug(
                    "sanitized_null_content",
                    tool_call_count=len(msg_copy["tool_calls"]),
                )

        sanitized.append(msg_copy)

    return sanitized


def get_ollama_provider() -> TatlockOllamaProvider:
    """
    Get a configured Ollama provider for PydanticAI agents.

    Returns:
        TatlockOllamaProvider configured with sanitization
    """
    return TatlockOllamaProvider()