"""
OpenAI model implementation (including OpenAI-compatible providers).
"""
import inspect
import json
import re
from typing import Any, AsyncIterator, Dict, List, Optional

import httpx
from openai import AsyncOpenAI

from ..base import BaseAIModel, Message, ModelConfig
from src.utils.logger import setup_logger

logger = setup_logger("OpenAIModel")


class OpenAIModel(BaseAIModel):
    """OpenAI model implementation."""

    def __init__(self, config: ModelConfig):
        super().__init__(config)
        self.logger = logger
        self._embedding_token_limit: Optional[int] = None

        http_client = httpx.AsyncClient(
            timeout=config.timeout,
            limits=httpx.Limits(max_keepalive_connections=5, max_connections=10),
        )

        self.client = AsyncOpenAI(
            api_key=config.api_key,
            base_url=config.api_base,
            timeout=config.timeout,
            http_client=http_client,
        )

        self._supports_tools = False
        self._supports_functions = False
        self._detect_tool_capability()

    def _detect_tool_capability(self) -> None:
        """Detect tool-calling parameters supported by current SDK."""
        try:
            signature = inspect.signature(self.client.chat.completions.create)
            parameters = signature.parameters
            self._supports_tools = "tools" in parameters
            self._supports_functions = "functions" in parameters
        except Exception as exc:
            self.logger.warning(f"Failed to inspect OpenAI completion signature: {exc}")

    def _build_tool_params(self, tools: Optional[List[dict]]) -> Dict[str, Any]:
        """Build SDK-compatible tool/function request parameters."""
        if not tools:
            return {}

        if self._supports_tools:
            return {"tools": tools}

        if self._supports_functions:
            functions = []
            for tool in tools:
                schema = self._extract_function_schema(tool)
                if schema:
                    functions.append(schema)

            if functions:
                return {"functions": functions, "function_call": "auto"}

        self.logger.warning(
            "Tool calling is not supported by current OpenAI SDK; tools were ignored."
        )
        return {}

    @staticmethod
    def _extract_function_schema(tool: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        if not isinstance(tool, dict):
            return None

        function_data = tool.get("function")
        if not isinstance(function_data, dict):
            return None

        name = function_data.get("name")
        if not name:
            return None

        return {
            "name": name,
            "description": function_data.get("description", ""),
            "parameters": function_data.get(
                "parameters", {"type": "object", "properties": {}}
            ),
        }

    async def _create_completion_with_fallback(
        self, params: Dict[str, Any], tools: Optional[List[dict]]
    ):
        """Create completion with runtime fallback for old SDK signatures."""
        try:
            return await self.client.chat.completions.create(**params)
        except TypeError as exc:
            message = str(exc)
            if "unexpected keyword argument 'tools'" in message and "tools" in params:
                self.logger.warning(
                    "SDK rejected `tools` at runtime, retrying with legacy functions."
                )
                self._supports_tools = False
                retry_params = dict(params)
                retry_params.pop("tools", None)
                retry_params.update(self._build_tool_params(tools))
                return await self.client.chat.completions.create(**retry_params)

            if "unexpected keyword argument 'functions'" in message and "functions" in params:
                self.logger.warning(
                    "SDK rejected `functions` at runtime, retrying without tool calling."
                )
                self._supports_functions = False
                retry_params = dict(params)
                retry_params.pop("functions", None)
                retry_params.pop("function_call", None)
                return await self.client.chat.completions.create(**retry_params)

            raise

    async def chat(
        self,
        messages: List[Message],
        tools: Optional[List[dict]] = None,
        **kwargs,
    ) -> Message:
        """Non-stream chat."""
        formatted_messages = [self._format_message(msg) for msg in messages]

        params = {
            "model": self.config.model_name,
            "messages": formatted_messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "frequency_penalty": self.config.frequency_penalty,
            "presence_penalty": self.config.presence_penalty,
        }

        params.update(self._build_tool_params(tools))
        params.update(kwargs)

        response = await self._create_completion_with_fallback(params, tools)

        choice = response.choices[0]
        tool_calls = self._extract_response_tool_calls(choice.message)
        return Message(
            role="assistant",
            content=choice.message.content or "",
            tool_calls=tool_calls,
        )

    async def chat_stream(
        self,
        messages: List[Message],
        tools: Optional[List[dict]] = None,
        **kwargs,
    ) -> AsyncIterator[str]:
        """Streaming chat."""
        formatted_messages = [self._format_message(msg) for msg in messages]

        params = {
            "model": self.config.model_name,
            "messages": formatted_messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "stream": True,
        }

        params.update(self._build_tool_params(tools))
        params.update(kwargs)

        stream = await self._create_completion_with_fallback(params, tools)

        async for chunk in stream:
            delta = chunk.choices[0].delta
            if delta.content:
                yield delta.content

    def _extract_response_tool_calls(self, message: Any) -> Optional[List[Dict[str, Any]]]:
        raw_tool_calls = getattr(message, "tool_calls", None)
        if raw_tool_calls:
            return [self._normalize_tool_call(tool_call) for tool_call in raw_tool_calls]

        legacy_function_call = getattr(message, "function_call", None)
        if legacy_function_call:
            return [self._normalize_legacy_function_call(legacy_function_call)]

        return None

    def _normalize_legacy_function_call(self, function_call: Any) -> Dict[str, Any]:
        if isinstance(function_call, dict):
            function_name = function_call.get("name")
            raw_arguments = function_call.get("arguments")
        else:
            function_name = getattr(function_call, "name", None)
            raw_arguments = getattr(function_call, "arguments", None)

        if isinstance(raw_arguments, dict):
            arguments = json.dumps(raw_arguments, ensure_ascii=False)
        elif raw_arguments is None:
            arguments = "{}"
        else:
            arguments = str(raw_arguments)

        return {
            "type": "function",
            "function": {
                "name": function_name,
                "arguments": arguments,
            },
        }

    def _format_message(self, msg: Message) -> Dict[str, Any]:
        """Convert internal message schema to OpenAI request format."""
        if msg.role == "assistant":
            formatted: Dict[str, Any] = {
                "role": "assistant",
                "content": msg.content if msg.content else None,
            }
            if msg.tool_calls:
                normalized_calls = [
                    self._normalize_tool_call(tool_call) for tool_call in msg.tool_calls
                ]
                if self._supports_tools:
                    formatted["tool_calls"] = normalized_calls
                elif self._supports_functions:
                    if len(normalized_calls) > 1:
                        self.logger.warning(
                            "Legacy function_call mode only supports one tool call; "
                            "extra tool calls were dropped."
                        )
                    function_data = normalized_calls[0].get("function") or {}
                    formatted["function_call"] = {
                        "name": function_data.get("name"),
                        "arguments": function_data.get("arguments", "{}"),
                    }
            return formatted

        if msg.role == "tool":
            if self._supports_tools and msg.tool_call_id:
                return {
                    "role": "tool",
                    "content": msg.content,
                    "tool_call_id": msg.tool_call_id,
                }

            return {
                "role": "function",
                "name": msg.name or "tool",
                "content": msg.content,
            }

        formatted = {"role": msg.role, "content": msg.content}
        if msg.name:
            formatted["name"] = msg.name
        return formatted

    def _normalize_tool_call(self, tool_call: Any) -> Dict[str, Any]:
        """Normalize tool call object/dict to a stable dict schema."""
        if isinstance(tool_call, dict):
            normalized = dict(tool_call)
        elif hasattr(tool_call, "model_dump"):
            normalized = tool_call.model_dump()
        else:
            function = getattr(tool_call, "function", None)
            if isinstance(function, dict):
                function_name = function.get("name")
                raw_arguments = function.get("arguments")
            else:
                function_name = getattr(function, "name", None)
                raw_arguments = getattr(function, "arguments", None)

            normalized = {
                "id": getattr(tool_call, "id", None),
                "type": getattr(tool_call, "type", "function"),
                "function": {
                    "name": function_name,
                    "arguments": raw_arguments,
                },
            }

        function_data = normalized.get("function") or {}
        if not isinstance(function_data, dict):
            function_data = {
                "name": getattr(function_data, "name", None),
                "arguments": getattr(function_data, "arguments", None),
            }
        raw_arguments = function_data.get("arguments")

        if isinstance(raw_arguments, dict):
            arguments = json.dumps(raw_arguments, ensure_ascii=False)
        elif raw_arguments is None:
            arguments = "{}"
        else:
            arguments = str(raw_arguments)

        function_data["arguments"] = arguments
        normalized["function"] = function_data
        normalized["type"] = normalized.get("type") or "function"

        return normalized

    @staticmethod
    def _is_embedding_too_long_error(error: Exception) -> bool:
        status_code = getattr(error, "status_code", None)
        if status_code == 413:
            return True

        message = str(error).lower()
        return (
            "less than 512 tokens" in message
            or "input too long" in message
            or "maximum context length" in message
        )

    @staticmethod
    def _extract_embedding_token_limit(error: Exception) -> Optional[int]:
        message = str(error).lower()
        patterns = [
            r"less than\s+(\d+)\s+tokens",
            r"maximum context length.*?(\d+)\s+tokens",
            r"max(?:imum)?(?: input)?(?: length)?\D+(\d+)\s+tokens",
        ]
        for pattern in patterns:
            match = re.search(pattern, message)
            if not match:
                continue

            try:
                limit = int(match.group(1))
            except (TypeError, ValueError):
                continue

            if limit > 0:
                return limit

        return None

    @staticmethod
    def _estimate_embedding_tokens(text: str) -> int:
        if not text:
            return 0

        ascii_count = 0
        cjk_count = 0
        other_count = 0

        for ch in text:
            code = ord(ch)
            if code < 128:
                ascii_count += 1
            elif 0x4E00 <= code <= 0x9FFF or 0x3400 <= code <= 0x4DBF:
                cjk_count += 1
            else:
                other_count += 1

        estimated = cjk_count + int(other_count * 0.7) + int(ascii_count / 4)
        return max(1, estimated)

    def _apply_embedding_budget(self, text: str) -> str:
        if not text or not self._embedding_token_limit:
            return text

        safe_limit = max(32, int(self._embedding_token_limit * 0.9))
        estimated_tokens = self._estimate_embedding_tokens(text)
        if estimated_tokens <= safe_limit:
            return text

        ratio = safe_limit / max(estimated_tokens, 1)
        target_length = max(64, int(len(text) * ratio))
        if target_length >= len(text):
            target_length = max(64, len(text) - 1)
        if target_length <= 0 or target_length >= len(text):
            return text

        head = target_length // 2
        tail = target_length - head
        return f"{text[:head]} {text[-tail:]}"

    @staticmethod
    def _shrink_text_for_embedding(text: str) -> str:
        compact = " ".join((text or "").split())
        if len(compact) <= 64:
            return compact

        target = max(64, int(len(compact) * 0.6))
        if target >= len(compact):
            target = max(64, len(compact) // 2)

        head = target // 2
        tail = target - head
        return f"{compact[:head]} {compact[-tail:]}"

    async def embed(self, text: str) -> List[float]:
        """Generate embeddings."""
        if isinstance(text, bytes):
            text = text.decode("utf-8", errors="ignore")

        raw_text = str(text or "")
        candidate_text = raw_text.strip() or raw_text or " "
        candidate_text = self._apply_embedding_budget(candidate_text)
        retry_count = 0

        while True:
            try:
                response = await self.client.embeddings.create(
                    model=self.config.model_name,
                    input=candidate_text,
                    encoding_format="float",
                )
                return response.data[0].embedding
            except Exception as e:
                if self._is_embedding_too_long_error(e):
                    token_limit = self._extract_embedding_token_limit(e)
                    if token_limit:
                        if not self._embedding_token_limit:
                            self._embedding_token_limit = token_limit
                        else:
                            self._embedding_token_limit = min(
                                self._embedding_token_limit, token_limit
                            )

                    next_text = self._apply_embedding_budget(
                        self._shrink_text_for_embedding(candidate_text)
                    )
                    if (
                        next_text
                        and len(next_text) < len(candidate_text)
                        and retry_count < 5
                    ):
                        retry_count += 1
                        limit_desc = (
                            f", token_limit={self._embedding_token_limit}"
                            if self._embedding_token_limit
                            else ""
                        )
                        self.logger.info(
                            "embedding input exceeded provider limit, retry with truncated text: "
                            f"{len(candidate_text)} -> {len(next_text)}{limit_desc}"
                        )
                        candidate_text = next_text
                        continue

                import traceback

                self.logger.error(f"embedding generation failed: {e}")
                self.logger.error(f"model: {self.config.model_name}")
                self.logger.error(f"text length: {len(candidate_text)}")
                self.logger.error(f"text preview: {repr(candidate_text[:100])}")
                self.logger.error(f"full traceback:\n{traceback.format_exc()}")
                raise
