"""OpenAI provider implementation."""

from typing import Any, AsyncIterator, cast

from openai import AsyncOpenAI, OpenAI
from openai.types.chat.chat_completion import ChatCompletion

from swarm.config.logging import get_structured_logger

from .base import ModelProvider
from .types import CompletionRequest

logger = get_structured_logger(__name__)


class OpenAIProvider(ModelProvider):
    """OpenAI provider for GPT models."""

    def __init__(
        self,
        sync_client: OpenAI,
        async_client: AsyncOpenAI,
    ) -> None:
        """Initialize OpenAI provider with injected clients.

        Args:
            sync_client: Configured OpenAI sync client
            async_client: Configured AsyncOpenAI client
        """
        self.sync_client = sync_client
        self.async_client = async_client

    def create_completion(
        self, request: CompletionRequest, model: str
    ) -> ChatCompletion:
        """Create a non-streaming completion."""
        params = self._prepare_params(request, stream=False, model=model)

        try:
            completion = cast(Any, self.sync_client).chat.completions.create(**params)
            return cast(ChatCompletion, completion)
        except Exception as e:
            logger.error(
                "openai_provider.completion_error",
                extra={
                    "location": "swarm.core.get_chat_completion",
                    "description": "OpenAI API error",
                    "error_type": type(e).__name__,
                    "error_message": str(e),
                    "status_code": getattr(e, "status_code", None),
                    "model": model,
                    "messages_count": len(request.messages),
                    "tools_count": len(request.tools) if request.tools else 0,
                    "function": "create_completion",
                },
            )
            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 cast(Any, self.async_client).chat.completions.create(**params)
        return cast(ChatCompletion, response)

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

        # Make the streaming API call
        try:
            stream = await cast(Any, self.async_client).chat.completions.create(
                **params
            )

            # Yield raw OpenAI chunks for now
            # DialogueRunner will handle conversion
            async for chunk in stream:
                yield chunk
        except Exception as e:
            logger.error(
                "openai_provider.stream_error",
                extra={
                    "error": str(e),
                    "model": model,
                },
            )
            raise

    def _prepare_params(
        self, request: CompletionRequest, stream: bool, model: str
    ) -> dict[str, Any]:
        """Prepare request parameters for OpenAI Chat Completions."""
        params: dict[str, Any] = {}
        # Extract model name after prefix (e.g., "openai/gpt-4o" -> "gpt-4o")
        actual_model = model.split("/", 1)[-1] if "/" in model else model
        params["model"] = actual_model
        params["messages"] = request.messages
        params["stream"] = stream
        if stream:
            params["stream_options"] = {"include_usage": True}

        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.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.response_format is not None:
            params["response_format"] = request.response_format
        if request.seed is not None:
            params["seed"] = request.seed

        return params
