"""OpenAI adapter (OpenAI chat/completions)."""

from __future__ import annotations

import httpx
import logging
from typing import Any

from mcpbridge.adapters._openai_style import append_tool_results_openai, extract_text_openai, extract_tool_calls_openai, is_done_openai
from mcpbridge.adapters.base import BaseLLMAdapter, ToolCall, ToolResult
from mcpbridge.config import LLMConfig
from mcpbridge.utils.token_counter import estimate_messages_tokens

logger = logging.getLogger("mcpbridge")


class OpenAIAdapter(BaseLLMAdapter):
    """Adapter for OpenAI's `chat/completions` endpoint."""

    def __init__(self, llm_config: LLMConfig) -> None:
        # Enforce provider dependency at init.
        try:
            import openai as _openai  # noqa: F401
        except ImportError as e:
            raise ImportError("Missing `openai` dependency. Install with: pip install mcpbridge[openai]") from e

        self._cfg = llm_config
        self._provider = "openai"

        self._model = llm_config.model
        self._api_key = llm_config.api_key
        self._base_url = (llm_config.base_url or "https://api.openai.com/v1").rstrip("/")
        self._timeout = int(llm_config.timeout)

        self._temperature = float(llm_config.temperature)
        self._max_tokens = int(llm_config.max_tokens)
        self._top_p = llm_config.top_p
        self._top_k = llm_config.top_k
        self._stop_sequences = llm_config.stop_sequences
        self._thinking = llm_config.thinking
        self._extra_params = dict(llm_config.extra_params or {})

    @property
    def provider_name(self) -> str:
        return self._provider

    @property
    def supports_parallel_tool_calls(self) -> bool:
        return True

    @property
    def supports_system_prompt(self) -> bool:
        return False

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        system: str | None = None,
        **kwargs: Any,
    ) -> Any:
        headers: dict[str, str] = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"

        if system and (not messages or messages[0].get("role") != "system"):
            messages = [{"role": "system", "content": system}, *messages]

        temperature = self._temperature
        if self._model.startswith(("o1", "o3")):
            temperature = None
            if messages and messages[0].get("role") == "system":
                sys_content = messages[0].get("content", "")
                messages = messages[1:]
                if messages and messages[0].get("role") == "user":
                    messages[0]["content"] = f"{sys_content}\n\n{messages[0].get('content','')}"
                else:
                    messages.insert(0, {"role": "user", "content": sys_content})

        payload: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "stream": False,
            "tools": tools if tools else None,
            "tool_choice": "auto" if tools else None,
            "max_tokens": self._max_tokens,
            **self._extra_params,
        }
        if temperature is not None:
            payload["temperature"] = temperature
        if self._top_p is not None:
            payload["top_p"] = self._top_p
        if self._stop_sequences:
            payload["stop"] = self._stop_sequences

        payload = {k: v for k, v in payload.items() if v is not None}

        url = f"{self._base_url}/chat/completions"
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, headers=headers, json=payload)
            resp.raise_for_status()
            return resp.json()

    def extract_tool_calls(self, response: Any) -> list[ToolCall]:
        return extract_tool_calls_openai(response)

    def extract_text(self, response: Any) -> str:
        return extract_text_openai(response)

    def is_done(self, response: Any) -> bool:
        return is_done_openai(response)

    def append_tool_results(self, messages: list[dict[str, Any]], response: Any, tool_results: list[ToolResult]) -> list[dict[str, Any]]:
        return append_tool_results_openai(messages, response, tool_results)

    def count_tokens(self, messages: list[dict[str, Any]], system: str) -> int:
        return estimate_messages_tokens(messages, system, provider=self._provider)

