"""
OpenAI 兼容 API 客户端

支持所有 OpenAI 兼容的 API：
- OpenAI (gpt-4, gpt-3.5-turbo 等)
- DeepSeek
- GLM (智谱)
- Qwen (通义千问)
- 其他兼容 OpenAI API 的服务
"""

import json
import httpx
from typing import AsyncGenerator, List, Dict, Any, Optional
from openai import AsyncOpenAI

from .base import (
    BaseModelClient,
    Message,
    MessageRole,
    ToolCall,
    ToolDefinition,
    StreamChunk,
    CompletionResponse,
    TokenUsage,
)
from core.logger import log_llm_request, log_llm_response, log_error, log_token_usage


class OpenAIClient(BaseModelClient):
    """
    OpenAI 兼容 API 客户端

    使用官方 openai SDK 实现
    """

    def __init__(
        self,
        model: str,
        api_key: str,
        base_url: Optional[str] = None,
        **kwargs
    ):
        """
        初始化 OpenAI 客户端

        Args:
            model: 模型名称
            api_key: API 密钥
            base_url: API 基础 URL（可选，用于兼容其他服务）
            **kwargs: 其他参数（temperature, max_tokens 等）
        """
        super().__init__(model, api_key, base_url, **kwargs)

        # 创建 AsyncOpenAI 客户端
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url,
            timeout=httpx.Timeout(300.0, connect=30.0),
            max_retries=5,
        )

    @property
    def provider(self) -> str:
        """返回提供商名称"""
        return "openai"

    def _convert_messages(self, messages: List[Message]) -> List[Dict[str, Any]]:
        """
        将统一消息格式转换为 OpenAI 格式

        Args:
            messages: 统一消息列表

        Returns:
            OpenAI 格式的消息列表
        """
        result = []
        for msg in messages:
            if msg.role == MessageRole.TOOL:
                # 工具结果消息
                if msg.tool_result:
                    result.append({
                        "role": "tool",
                        "tool_call_id": msg.tool_result.tool_call_id,
                        "content": msg.tool_result.content,
                    })
            elif msg.role == MessageRole.ASSISTANT and msg.tool_calls:
                # 带工具调用的助手消息
                tool_calls = []
                for tc in msg.tool_calls:
                    tool_calls.append({
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.name,
                            "arguments": json.dumps(tc.arguments, ensure_ascii=False),
                        }
                    })
                result.append({
                    "role": "assistant",
                    "content": msg.content or None,
                    "tool_calls": tool_calls,
                })
            else:
                # 普通消息
                result.append({
                    "role": msg.role.value,
                    "content": msg.content,
                })
        return result

    def _convert_tools(self, tools: List[ToolDefinition]) -> List[Dict[str, Any]]:
        """
        将工具定义转换为 OpenAI 格式

        Args:
            tools: 工具定义列表

        Returns:
            OpenAI 格式的工具列表
        """
        result = []
        for tool in tools:
            result.append({
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                }
            })
        return result

    def _parse_tool_calls(self, tool_calls: List[Any]) -> List[ToolCall]:
        """
        解析 OpenAI 返回的工具调用

        Args:
            tool_calls: OpenAI 格式的工具调用

        Returns:
            统一格式的工具调用列表
        """
        result = []
        for tc in tool_calls:
            try:
                arguments = json.loads(tc.function.arguments)
            except (json.JSONDecodeError, AttributeError):
                arguments = {}

            result.append(ToolCall(
                id=tc.id,
                name=tc.function.name,
                arguments=arguments,
            ))
        return result

    async def chat_stream(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> AsyncGenerator[StreamChunk, None]:
        """
        流式对话

        Args:
            messages: 消息列表
            tools: 工具定义列表
            system_prompts: 系统提示词列表
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Yields:
            StreamChunk: 流式输出块
        """
        # 三级优先级: 入参 > 实例属性 > 兜底
        effective_temperature = temperature if temperature is not None else self.temperature
        effective_max_tokens = max_tokens if max_tokens is not None else self.max_tokens

        # 转换消息格式
        openai_messages = self._convert_messages(messages)

        # 系统提示词：拼接成一条 system message 插入最前面
        if system_prompts:
            combined_system = "\n\n".join(m.content for m in system_prompts)
            openai_messages.insert(0, {"role": "system", "content": combined_system})

        # 构建请求参数
        request_params = {
            "model": self.model,
            "messages": openai_messages,
            "temperature": effective_temperature,
            "stream": True,
            "stream_options": {"include_usage": True},
        }

        if effective_max_tokens:
            request_params["max_tokens"] = effective_max_tokens

        if tools:
            request_params["tools"] = self._convert_tools(tools)
            request_params["tool_choice"] = "auto"

        # 合并额外参数
        request_params.update(kwargs)

        log_llm_request(
            provider=self.provider,
            model=self.model,
            request_body=request_params,
        )

        # 用于累积工具调用
        tool_call_chunks: Dict[int, Dict[str, str]] = {}
        # 用于累积文本内容（用于日志记录）
        accumulated_text: str = ""

        try:
            response = await self.client.chat.completions.create(**request_params)

            final_usage: Optional[TokenUsage] = None
            finish_reason: Optional[str] = None
            text_completed = False

            async for chunk in response:
                # usage-only chunk: choices 为空（OpenAI 在流末尾单独发送）
                if hasattr(chunk, "usage") and chunk.usage:
                    prompt = chunk.usage.prompt_tokens or 0
                    cached = 0
                    if hasattr(chunk.usage, 'prompt_tokens_details') and chunk.usage.prompt_tokens_details:
                        cached = getattr(chunk.usage.prompt_tokens_details, 'cached_tokens', 0) or 0
                    final_usage = TokenUsage(
                        input_tokens=prompt - cached,
                        output_tokens=chunk.usage.completion_tokens or 0,
                        cache_read_input_tokens=cached,
                    )

                delta = chunk.choices[0].delta
                choice_finish = chunk.choices[0].finish_reason

                stream_chunk = StreamChunk()

                # 处理文本内容
                if delta and delta.content:
                    stream_chunk.text = delta.content
                    accumulated_text += delta.content

                # 处理工具调用（流式累积）
                if delta and delta.tool_calls:
                    # 首次收到 tool_calls delta → 文本流已结束
                    if not text_completed and accumulated_text:
                        text_completed = True
                        yield StreamChunk(text_complete=True)

                    for tc_delta in delta.tool_calls:
                        idx = tc_delta.index
                        if idx not in tool_call_chunks:
                            tool_call_chunks[idx] = {
                                "id": "",
                                "name": "",
                                "arguments": "",
                            }

                        # id/name 首次赋值，后续忽略（防御重复下发）
                        if tc_delta.id and not tool_call_chunks[idx]["id"]:
                            tool_call_chunks[idx]["id"] = tc_delta.id
                        if tc_delta.function:
                            if tc_delta.function.name and not tool_call_chunks[idx]["name"]:
                                tool_call_chunks[idx]["name"] = tc_delta.function.name
                            if tc_delta.function.arguments:
                                tool_call_chunks[idx]["arguments"] += tc_delta.function.arguments

                # 记录 finish_reason（但不在这里 yield is_finished）
                if choice_finish:
                    finish_reason = choice_finish
                    # 纯文本响应结束（无工具调用）
                    if not text_completed and accumulated_text and not tool_call_chunks:
                        text_completed = True
                        yield StreamChunk(text_complete=True)

                # 有文本内容时 yield
                if stream_chunk.text:
                    yield stream_chunk

            # --- 流结束，统一 yield 最终 chunk ---
            final_chunk = StreamChunk(
                is_finished=True,
                finish_reason=finish_reason,
            )

            # 解析累积的工具调用
            if tool_call_chunks:
                for idx in sorted(tool_call_chunks.keys()):
                    tc = tool_call_chunks[idx]
                    try:
                        arguments = json.loads(tc["arguments"]) if tc["arguments"] else {}
                    except json.JSONDecodeError:
                        arguments = {}

                    final_chunk.tool_calls.append(ToolCall(
                        id=tc["id"],
                        name=tc["name"],
                        arguments=arguments,
                    ))

            final_chunk.usage = final_usage

            # 日志
            if final_usage:
                log_token_usage(
                    provider="openai",
                    model=self.model,
                    input_tokens=final_usage.input_tokens,
                    output_tokens=final_usage.output_tokens,
                    cache_read_input_tokens=final_usage.cache_read_input_tokens,
                )

            log_llm_response(
                provider=self.provider,
                model=self.model,
                response_body={
                    "content": accumulated_text,
                    "tool_calls": [tc.__dict__ for tc in final_chunk.tool_calls] if final_chunk.tool_calls else None,
                    "usage": final_chunk.usage.to_dict() if final_chunk.usage else None,
                },
            )

            yield final_chunk

        except Exception as e:
            import traceback
            log_error(
                f"openai_stream_error: [{type(e).__name__}] {str(e) or repr(e)}\n{traceback.format_exc()}",
                model=self.model,
            )
            raise

    async def chat(
        self,
        messages: List[Message],
        tools: Optional[List[ToolDefinition]] = None,
        system_prompts: Optional[List[Message]] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> CompletionResponse:
        """
        非流式对话

        Args:
            messages: 消息列表
            tools: 工具定义列表
            system_prompts: 系统提示词列表
            temperature: 温度参数（None 时使用实例默认值）
            max_tokens: 最大 token 数（None 时使用实例默认值）
            **kwargs: 其他参数

        Returns:
            CompletionResponse: 完整响应
        """
        # 三级优先级: 入参 > 实例属性 > 兜底
        effective_temperature = temperature if temperature is not None else self.temperature
        effective_max_tokens = max_tokens if max_tokens is not None else self.max_tokens

        # 转换消息格式
        openai_messages = self._convert_messages(messages)

        # 系统提示词：拼接成一条 system message 插入最前面
        if system_prompts:
            combined_system = "\n\n".join(m.content for m in system_prompts)
            openai_messages.insert(0, {"role": "system", "content": combined_system})

        # 构建请求参数
        request_params = {
            "model": self.model,
            "messages": openai_messages,
            "temperature": effective_temperature,
        }

        if effective_max_tokens:
            request_params["max_tokens"] = effective_max_tokens

        if tools:
            request_params["tools"] = self._convert_tools(tools)
            request_params["tool_choice"] = "auto"

        # 合并额外参数
        request_params.update(kwargs)

        log_llm_request(
            provider=self.provider,
            model=self.model,
            request_body=request_params,
        )

        try:
            response = await self.client.chat.completions.create(**request_params)

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

            result = CompletionResponse(
                content=message.content or "",
                finish_reason=choice.finish_reason,
            )

            # 解析工具调用
            if message.tool_calls:
                result.tool_calls = self._parse_tool_calls(message.tool_calls)

            # 获取 usage（转换为统一的 TokenUsage 格式）
            if response.usage:
                prompt = response.usage.prompt_tokens or 0
                cached = 0
                if hasattr(response.usage, 'prompt_tokens_details') and response.usage.prompt_tokens_details:
                    cached = getattr(response.usage.prompt_tokens_details, 'cached_tokens', 0) or 0
                result.usage = TokenUsage(
                    input_tokens=prompt - cached,
                    output_tokens=response.usage.completion_tokens or 0,
                    cache_read_input_tokens=cached,
                )

                log_token_usage(
                    provider="openai",
                    model=self.model,
                    input_tokens=result.usage.input_tokens,
                    output_tokens=result.usage.output_tokens,
                    cache_read_input_tokens=result.usage.cache_read_input_tokens,
                )

            log_llm_response(
                provider=self.provider,
                model=self.model,
                response_body={
                    "content": result.content,
                    "tool_calls": [tc.__dict__ for tc in result.tool_calls] if result.tool_calls else None,
                    "usage": result.usage.to_dict() if result.usage else None,
                },
            )

            return result

        except Exception as e:
            import traceback
            log_error(
                f"openai_chat_error: [{type(e).__name__}] {str(e) or repr(e)}\n{traceback.format_exc()}",
                model=self.model,
            )
            raise

    async def count_tokens(
        self,
        messages: List[Message],
        system_prompts: Optional[List[Message]] = None,
        tools: Optional[List[ToolDefinition]] = None,
    ) -> int:
        """
        使用 tiktoken 本地计算 token 数

        失败时 fallback 到基类的字符估算。
        """
        try:
            import tiktoken
        except ImportError:
            return await super().count_tokens(messages, system_prompts, tools)

        try:
            try:
                encoding = tiktoken.encoding_for_model(self.model)
            except KeyError:
                encoding = tiktoken.get_encoding("cl100k_base")

            num_tokens = 0

            # 系统提示词
            if system_prompts:
                combined = "\n\n".join(m.content for m in system_prompts)
                num_tokens += len(encoding.encode(combined)) + 3  # system 消息开销

            # 对话消息
            for msg in messages:
                num_tokens += 3  # 每条消息固定开销
                if msg.content:
                    num_tokens += len(encoding.encode(msg.content))
                for tc in (msg.tool_calls or []):
                    num_tokens += len(encoding.encode(tc.name))
                    num_tokens += len(encoding.encode(json.dumps(tc.arguments, ensure_ascii=False)))
                if msg.tool_result:
                    num_tokens += len(encoding.encode(msg.tool_result.content))

            # 工具定义
            if tools:
                for tool in tools:
                    num_tokens += len(encoding.encode(tool.name))
                    num_tokens += len(encoding.encode(tool.description))
                    num_tokens += len(encoding.encode(json.dumps(tool.parameters, ensure_ascii=False)))

            num_tokens += 3  # 回复前缀开销
            return num_tokens

        except Exception as e:
            log_error(f"tiktoken_count_error: {e}")
            return await super().count_tokens(messages, system_prompts, tools)
