"""OpenAI model provider — implements ModelProvider using the OpenAI SDK."""

from __future__ import annotations

import json

from hatchdx.agent.providers.base import (
    AgentResponse,
    ProviderAuthError,
    ProviderError,
    ToolCall,
)


# ---------------------------------------------------------------------------
# Pricing (per million tokens)
# ---------------------------------------------------------------------------

_PRICING: dict[str, dict[str, float]] = {
    # GPT-5 family
    "gpt-5.2": {"input": 1.75, "output": 14.0},
    "gpt-5.1": {"input": 1.25, "output": 10.0},
    "gpt-5": {"input": 1.25, "output": 10.0},
    "gpt-5-mini": {"input": 0.25, "output": 2.0},
    "gpt-5-nano": {"input": 0.05, "output": 0.40},
    # GPT-4.1 family
    "gpt-4.1": {"input": 2.0, "output": 8.0},
    "gpt-4.1-mini": {"input": 0.40, "output": 1.60},
    "gpt-4.1-nano": {"input": 0.10, "output": 0.40},
    # GPT-4o family (legacy)
    "gpt-4o": {"input": 2.50, "output": 10.0},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    # Reasoning models
    "o3": {"input": 2.0, "output": 8.0},
    "o4-mini": {"input": 1.10, "output": 4.40},
    "o1": {"input": 15.0, "output": 60.0},
    "o3-mini": {"input": 1.10, "output": 4.40},
    # Legacy
    "gpt-4-turbo": {"input": 10.0, "output": 30.0},
    "gpt-4": {"input": 30.0, "output": 60.0},
    "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
}

_DEFAULT_PRICING = {"input": 1.25, "output": 10.0}


def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Estimate cost in USD for a given OpenAI model and token counts."""
    pricing = _PRICING.get(model, _DEFAULT_PRICING)
    return (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000


# ---------------------------------------------------------------------------
# Tool format conversion
# ---------------------------------------------------------------------------


def mcp_tools_to_openai(tools: list[dict]) -> list[dict]:
    """Convert MCP tool definitions to OpenAI function calling format.

    MCP format:
        {"name": "...", "description": "...", "inputSchema": {...}}

    OpenAI format:
        {"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}
    """
    result = []
    for tool in tools:
        fn: dict = {
            "name": tool["name"],
            "description": tool.get("description", ""),
        }
        schema = tool.get("inputSchema") or tool.get("input_schema")
        if schema:
            fn["parameters"] = schema
        else:
            fn["parameters"] = {"type": "object", "properties": {}}
        result.append({"type": "function", "function": fn})
    return result


def _openai_messages(
    messages: list[dict], system: str
) -> list[dict]:
    """Convert internal message format to OpenAI chat messages.

    Prepends the system prompt as a system message and maps tool results
    to OpenAI's expected format.
    """
    result: list[dict] = [{"role": "system", "content": system}]

    for msg in messages:
        role = msg.get("role")

        if role == "user":
            result.append({"role": "user", "content": msg["content"]})

        elif role == "assistant":
            entry: dict = {"role": "assistant"}
            if msg.get("content"):
                entry["content"] = msg["content"]
            if msg.get("tool_calls"):
                entry["tool_calls"] = [
                    {
                        "id": tc.id if hasattr(tc, "id") else tc["id"],
                        "type": "function",
                        "function": {
                            "name": tc.name if hasattr(tc, "name") else tc["name"],
                            "arguments": json.dumps(
                                tc.arguments if hasattr(tc, "arguments") else tc["arguments"]
                            ),
                        },
                    }
                    for tc in msg["tool_calls"]
                ]
            result.append(entry)

        elif role == "tool":
            result.append({
                "role": "tool",
                "tool_call_id": msg["tool_call_id"],
                "content": msg["content"],
            })

    return result


# ---------------------------------------------------------------------------
# Provider
# ---------------------------------------------------------------------------


class OpenAIProvider:
    """OpenAI model provider using the openai SDK."""

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        base_url: str | None = None,
    ) -> None:
        self.model = model
        self._api_key = api_key
        self._base_url = base_url
        self._client = None

    def _get_client(self):
        """Lazy-init the OpenAI client."""
        if self._client is not None:
            return self._client

        try:
            import openai
        except ImportError:
            raise ProviderError(
                "OpenAI SDK not installed.\n\n"
                "Install it with:\n"
                "  uv add hatchdx --extra openai"
            )

        import os

        api_key = self._api_key or os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise ProviderAuthError(
                "Missing OpenAI API key.\n\n"
                "Set it with: export OPENAI_API_KEY=your-key-here\n"
                "Or pass it in agent.yaml model config."
            )

        kwargs: dict = {"api_key": api_key}
        if self._base_url:
            kwargs["base_url"] = self._base_url

        self._client = openai.AsyncOpenAI(**kwargs)
        return self._client

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict],
        system: str,
        max_tokens: int,
        temperature: float,
    ) -> AgentResponse:
        """Send a chat request to OpenAI and return an AgentResponse."""
        client = self._get_client()
        openai_messages = _openai_messages(messages, system)
        openai_tools = mcp_tools_to_openai(tools) if tools else None

        kwargs: dict = {
            "model": self.model,
            "messages": openai_messages,
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        if openai_tools:
            kwargs["tools"] = openai_tools

        try:
            import openai as openai_module

            response = await client.chat.completions.create(**kwargs)
        except openai_module.AuthenticationError as e:
            raise ProviderAuthError(f"OpenAI authentication failed: {e}") from e
        except openai_module.APIError as e:
            raise ProviderError(f"OpenAI API error: {e}") from e

        choice = response.choices[0]
        message = choice.message

        # Extract text
        text = message.content

        # Extract tool calls
        tool_calls = []
        if message.tool_calls:
            for tc in message.tool_calls:
                try:
                    arguments = json.loads(tc.function.arguments)
                except (json.JSONDecodeError, TypeError):
                    arguments = {}
                tool_calls.append(
                    ToolCall(
                        id=tc.id,
                        name=tc.function.name,
                        arguments=arguments,
                    )
                )

        # Map finish reason
        stop_reason = "end_turn"
        if choice.finish_reason == "tool_calls":
            stop_reason = "tool_use"
        elif choice.finish_reason == "length":
            stop_reason = "max_tokens"

        # Usage
        usage = {"input": 0, "output": 0}
        if response.usage:
            usage["input"] = response.usage.prompt_tokens or 0
            usage["output"] = response.usage.completion_tokens or 0

        return AgentResponse(
            text=text,
            tool_calls=tool_calls,
            stop_reason=stop_reason,
            usage=usage,
        )
