"""LiteLLM provider for multi-model support."""

from typing import Any, AsyncIterator, Callable, Dict, cast

from litellm import acompletion, completion
from openai.types.chat import ChatCompletion

from swarm.config.logging import get_structured_logger

from .base import ModelProvider
from .types import CompletionRequest

logger = get_structured_logger(__name__)


class LiteLLMProvider(ModelProvider):
    """LiteLLM provider supporting multiple model providers."""

    def __init__(self, config_getter: Callable[[], Dict[str, str]]) -> None:
        """Initialize LiteLLM provider with config getter.

        Args:
            config_getter: Function that returns provider-specific config
        """
        self.config_getter = config_getter

    def _prepare_params(
        self, request: CompletionRequest, stream: bool, model: str
    ) -> Dict[str, Any]:
        """Prepare parameters for LiteLLM call."""
        params: Dict[str, Any] = {
            "model": model,
            "messages": request.messages,
            "stream": stream,
        }
        if stream:
            params["stream_options"] = {"include_usage": True}

        # Add optional parameters
        if request.temperature is not None:
            params["temperature"] = request.temperature
        if request.max_tokens is not None:
            params["max_tokens"] = request.max_tokens
        if request.top_p is not None:
            params["top_p"] = request.top_p
        if request.frequency_penalty is not None:
            params["frequency_penalty"] = request.frequency_penalty
        if request.presence_penalty is not None:
            params["presence_penalty"] = request.presence_penalty
        if request.stop is not None:
            params["stop"] = request.stop
        if request.tools is not None:
            params["tools"] = request.tools
        if request.tool_choice is not None:
            params["tool_choice"] = request.tool_choice
        if request.response_format is not None:
            params["response_format"] = request.response_format
        if request.seed is not None:
            params["seed"] = request.seed

        # Merge provider-specific config
        params.update(self.config_getter())

        return params

    def create_completion(
        self, request: CompletionRequest, model: str
    ) -> ChatCompletion:
        """Create a non-streaming completion."""
        try:
            params = self._prepare_params(request, stream=False, model=model)
            response = completion(**params)
            # LiteLLM returns OpenAI-compatible responses
            # Cast is safe as the structure matches ChatCompletion
            return cast(ChatCompletion, response)
        except Exception as e:
            logger.error(
                "litellm_provider.completion_error",
                extra={
                    "error": str(e),
                    "model": model,
                },
            )
            raise

    async def create_completion_async(
        self, request: CompletionRequest, model: str
    ) -> ChatCompletion:
        """Create a non-streaming completion (async)."""
        params = self._prepare_params(request, stream=False, model=model)
        response = await acompletion(**params)
        # Cast ModelResponse to ChatCompletion (they're compatible)
        return cast(ChatCompletion, response)

    async def stream_completion(
        self, request: CompletionRequest, model: str
    ) -> AsyncIterator[Any]:
        """Create a streaming completion."""
        try:
            params = self._prepare_params(request, stream=True, model=model)

            stream = await acompletion(**params)
            async for chunk in stream:
                yield chunk
        except Exception as e:
            logger.error(
                "litellm_provider.stream_error",
                extra={
                    "error": str(e),
                    "model": model,
                },
            )
            raise
