"""
OpenAI-compatible model provider.

Supports:
  - OpenAI GPT-4o, GPT-4, GPT-3.5-turbo, o1, o3 etc.
  - Any OpenAI-compatible API (Azure, LiteLLM, vLLM, etc.)
  - Function/tool calling with parallel tool calls
  - Streaming (future)
"""

from __future__ import annotations

import os
from typing import Any

from nexus.core.message import Message
from nexus.models.base import ModelProvider
from nexus.types import TokenUsage


class OpenAIProvider(ModelProvider):
    """
    LLM provider for OpenAI and compatible APIs.

    Usage:
        provider = OpenAIProvider(model="gpt-4o")
        response = await provider.chat(messages=[...])
    """

    def __init__(
        self,
        model: str = "gpt-4o",
        api_key: str | None = None,
        base_url: str | None = None,
        organization: str | None = None,
        default_temperature: float = 0.7,
        default_max_tokens: int | None = None,
    ) -> None:
        self._model = model
        self._api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
        self._base_url = base_url or os.environ.get(
            "OPENAI_BASE_URL", "https://api.openai.com/v1"
        )
        self._organization = organization or os.environ.get("OPENAI_ORG_ID")
        self._default_temperature = default_temperature
        self._default_max_tokens = default_max_tokens
        self._client: Any = None  # lazy init

    async def _get_client(self) -> Any:
        """Lazy-initialize the openai async client."""
        if self._client is None:
            try:
                from openai import AsyncOpenAI
            except ImportError:
                raise ImportError(
                    "openai package is required for OpenAIProvider. "
                    "Install it with: pip install nexus-agents[openai]"
                )
            self._client = AsyncOpenAI(
                api_key=self._api_key,
                base_url=self._base_url,
                organization=self._organization,
            )
        return self._client

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        client = await self._get_client()

        # Build request params
        params: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "temperature": temperature if temperature is not None else self._default_temperature,
        }

        effective_max_tokens = max_tokens or self._default_max_tokens
        if effective_max_tokens:
            params["max_tokens"] = effective_max_tokens

        if tools:
            params["tools"] = tools
            params["tool_choice"] = kwargs.pop("tool_choice", "auto")

        # Merge extra kwargs
        params.update(kwargs)

        # Call API
        response = await client.chat.completions.create(**params)
        choice = response.choices[0]
        resp_message = choice.message

        # Extract tool calls
        tool_calls = None
        if resp_message.tool_calls:
            tool_calls = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {
                        "name": tc.function.name,
                        "arguments": tc.function.arguments,
                    },
                }
                for tc in resp_message.tool_calls
            ]

        # Build normalized Message
        message = Message.agent(
            content=resp_message.content or "",
            agent_name=self._model,
        )
        # Attach tool calls as metadata
        if tool_calls:
            message.metadata["tool_calls"] = tool_calls

        # Build token usage
        usage = TokenUsage(
            prompt_tokens=response.usage.prompt_tokens if response.usage else 0,
            completion_tokens=response.usage.completion_tokens if response.usage else 0,
            total_tokens=response.usage.total_tokens if response.usage else 0,
        )

        return {
            "message": message,
            "usage": usage,
            "model": response.model,
            "finish_reason": choice.finish_reason,
            "tool_calls": tool_calls,
            "raw": response,
        }

    async def close(self) -> None:
        if self._client is not None:
            await self._client.close()
            self._client = None

    @property
    def model_name(self) -> str:
        return self._model
