"""
OpenAI 提供商适配器
"""
import logging
from typing import AsyncGenerator, Dict, List, Any

from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
from .registry import ProviderRegistry

logger = logging.getLogger(__name__)


@ProviderRegistry.register
class OpenAIProvider(BaseLLMProvider):
    """
    OpenAI 提供商适配器
    
    支持 GPT-3.5, GPT-4 等模型
    """
    
    provider_type = 'openai'
    provider_name = 'OpenAI'
    supported_model_types = ['chat', 'completion', 'embedding']
    
    DEFAULT_API_BASE = 'https://api.openai.com/v1'
    
    def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
        super().__init__(api_key, api_base, **kwargs)
        self.api_base = api_base or self.DEFAULT_API_BASE
        self._client = None
        self._async_client = None
    
    def _get_client(self):
        """获取同步客户端"""
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI(
                    api_key=self.api_key,
                    base_url=self.api_base,
                )
            except ImportError:
                raise ImportError('请安装 openai 库: pip install openai')
        return self._client
    
    def _get_async_client(self):
        """获取异步客户端"""
        if self._async_client is None:
            try:
                from openai import AsyncOpenAI
                self._async_client = AsyncOpenAI(
                    api_key=self.api_key,
                    base_url=self.api_base,
                )
            except ImportError:
                raise ImportError('请安装 openai 库: pip install openai')
        return self._async_client
    
    def chat(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> LLMResponse:
        """同步对话"""
        client = self._get_client()
        
        # 构建请求参数
        kwargs = {
            'model': config.model,
            'messages': [m.to_dict() for m in messages],
            'temperature': config.temperature,
            'top_p': config.top_p,
            'max_tokens': config.max_tokens,
            'stop': config.stop,
            'presence_penalty': config.presence_penalty,
            'frequency_penalty': config.frequency_penalty,
            **config.extra_params,
        }
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_openai_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = 'required'
            elif config.tool_choice == 'none':
                kwargs['tool_choice'] = 'none'
            elif config.tool_choice != 'auto':
                # 指定具体工具
                kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
            else:
                kwargs['tool_choice'] = 'auto'
        
        response = client.chat.completions.create(**kwargs)
        
        choice = response.choices[0]
        usage = response.usage
        
        # 解析工具调用
        tool_calls = None
        if choice.message.tool_calls:
            tool_calls = self._parse_tool_calls(choice.message.tool_calls)
        
        return LLMResponse(
            content=choice.message.content or '',
            model=response.model,
            prompt_tokens=usage.prompt_tokens if usage else 0,
            completion_tokens=usage.completion_tokens if usage else 0,
            total_tokens=usage.total_tokens if usage else 0,
            finish_reason=choice.finish_reason or '',
            raw_response=response.model_dump(),
            tool_calls=tool_calls,
        )
    
    def _parse_tool_calls(self, tool_calls) -> List[ToolCall]:
        """解析工具调用"""
        import json
        result = []
        for tc in tool_calls:
            try:
                arguments = json.loads(tc.function.arguments) if tc.function.arguments else {}
            except json.JSONDecodeError:
                arguments = {'raw': tc.function.arguments}
            
            result.append(ToolCall(
                id=tc.id,
                name=tc.function.name,
                arguments=arguments,
            ))
        return result
    
    async def chat_async(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> LLMResponse:
        """异步对话"""
        client = self._get_async_client()
        
        # 构建请求参数
        kwargs = {
            'model': config.model,
            'messages': [m.to_dict() for m in messages],
            'temperature': config.temperature,
            'top_p': config.top_p,
            'max_tokens': config.max_tokens,
            'stop': config.stop,
            'presence_penalty': config.presence_penalty,
            'frequency_penalty': config.frequency_penalty,
            **config.extra_params,
        }
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_openai_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = 'required'
            elif config.tool_choice == 'none':
                kwargs['tool_choice'] = 'none'
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
            else:
                kwargs['tool_choice'] = 'auto'
        
        response = await client.chat.completions.create(**kwargs)
        
        choice = response.choices[0]
        usage = response.usage
        
        # 解析工具调用
        tool_calls = None
        if choice.message.tool_calls:
            tool_calls = self._parse_tool_calls(choice.message.tool_calls)
        
        return LLMResponse(
            content=choice.message.content or '',
            model=response.model,
            prompt_tokens=usage.prompt_tokens if usage else 0,
            completion_tokens=usage.completion_tokens if usage else 0,
            total_tokens=usage.total_tokens if usage else 0,
            finish_reason=choice.finish_reason or '',
            raw_response=response.model_dump(),
            tool_calls=tool_calls,
        )
    
    async def chat_stream(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> AsyncGenerator[LLMStreamChunk, None]:
        """流式对话"""
        client = self._get_async_client()
        
        # 构建请求参数
        kwargs = {
            'model': config.model,
            'messages': [m.to_dict() for m in messages],
            'temperature': config.temperature,
            'top_p': config.top_p,
            'max_tokens': config.max_tokens,
            'stop': config.stop,
            'presence_penalty': config.presence_penalty,
            'frequency_penalty': config.frequency_penalty,
            'stream': True,
            'stream_options': {"include_usage": True},
            **config.extra_params,
        }
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_openai_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = 'required'
            elif config.tool_choice == 'none':
                kwargs['tool_choice'] = 'none'
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
            else:
                kwargs['tool_choice'] = 'auto'
        
        stream = await client.chat.completions.create(**kwargs)
        
        # 用于累积工具调用
        tool_call_accumulator = {}  # id -> {name, arguments}
        
        async for chunk in stream:
            if chunk.choices:
                choice = chunk.choices[0]
                delta = choice.delta
                
                # 处理工具调用增量
                tool_call_delta = None
                if hasattr(delta, 'tool_calls') and delta.tool_calls:
                    for tc_delta in delta.tool_calls:
                        tc_index = tc_delta.index
                        if tc_index not in tool_call_accumulator:
                            tool_call_accumulator[tc_index] = {
                                'id': tc_delta.id or '',
                                'name': '',
                                'arguments': '',
                            }
                        if tc_delta.id:
                            tool_call_accumulator[tc_index]['id'] = tc_delta.id
                        if tc_delta.function:
                            if tc_delta.function.name:
                                tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
                            if tc_delta.function.arguments:
                                tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
                        tool_call_delta = tool_call_accumulator[tc_index]
                
                # 检查是否完成（包含工具调用）
                tool_calls = None
                if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
                    tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
                
                yield LLMStreamChunk(
                    content=delta.content or '',
                    is_finished=choice.finish_reason is not None,
                    finish_reason=choice.finish_reason or '',
                    tool_calls=tool_calls,
                    tool_call_delta=tool_call_delta,
                )
            
            # 最后一个 chunk 包含 usage 信息
            if chunk.usage:
                yield LLMStreamChunk(
                    content='',
                    is_finished=True,
                    prompt_tokens=chunk.usage.prompt_tokens,
                    completion_tokens=chunk.usage.completion_tokens,
                    total_tokens=chunk.usage.total_tokens,
                )
    
    def _parse_accumulated_tool_calls(self, accumulator: Dict) -> List[ToolCall]:
        """解析累积的工具调用"""
        import json
        result = []
        for idx in sorted(accumulator.keys()):
            tc = accumulator[idx]
            try:
                arguments = json.loads(tc['arguments']) if tc['arguments'] else {}
            except json.JSONDecodeError:
                arguments = {'raw': tc['arguments']}
            
            result.append(ToolCall(
                id=tc['id'],
                name=tc['name'],
                arguments=arguments,
            ))
        return result
    
    def chat_stream_sync(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ):
        """同步流式对话"""
        client = self._get_client()
        
        # 构建请求参数
        kwargs = {
            'model': config.model,
            'messages': [m.to_dict() for m in messages],
            'temperature': config.temperature,
            'top_p': config.top_p,
            'max_tokens': config.max_tokens,
            'stop': config.stop,
            'presence_penalty': config.presence_penalty,
            'frequency_penalty': config.frequency_penalty,
            'stream': True,
            'stream_options': {"include_usage": True},
            **config.extra_params,
        }
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_openai_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = 'required'
            elif config.tool_choice == 'none':
                kwargs['tool_choice'] = 'none'
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
            else:
                kwargs['tool_choice'] = 'auto'
        
        stream = client.chat.completions.create(**kwargs)
        
        # 用于累积工具调用
        tool_call_accumulator = {}
        
        for chunk in stream:
            if chunk.choices:
                choice = chunk.choices[0]
                delta = choice.delta
                
                # 处理工具调用增量
                tool_call_delta = None
                if hasattr(delta, 'tool_calls') and delta.tool_calls:
                    for tc_delta in delta.tool_calls:
                        tc_index = tc_delta.index
                        if tc_index not in tool_call_accumulator:
                            tool_call_accumulator[tc_index] = {
                                'id': tc_delta.id or '',
                                'name': '',
                                'arguments': '',
                            }
                        if tc_delta.id:
                            tool_call_accumulator[tc_index]['id'] = tc_delta.id
                        if tc_delta.function:
                            if tc_delta.function.name:
                                tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
                            if tc_delta.function.arguments:
                                tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
                        tool_call_delta = tool_call_accumulator[tc_index]
                
                # 检查是否完成（包含工具调用）
                tool_calls = None
                if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
                    tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
                
                yield LLMStreamChunk(
                    content=delta.content or '',
                    is_finished=choice.finish_reason is not None,
                    finish_reason=choice.finish_reason or '',
                    tool_calls=tool_calls,
                    tool_call_delta=tool_call_delta,
                )
            
            # 最后一个 chunk 包含 usage 信息
            if chunk.usage:
                yield LLMStreamChunk(
                    content='',
                    is_finished=True,
                    prompt_tokens=chunk.usage.prompt_tokens,
                    completion_tokens=chunk.usage.completion_tokens,
                    total_tokens=chunk.usage.total_tokens,
                )
    
    # 模型类型推断关键词
    _EMBEDDING_KEYWORDS = ('embed', 'bge', 'nomic-embed', 'mxbai-embed', 'jina-embedding')
    _RERANK_KEYWORDS = ('rerank', 'reranker')
    _SKIP_KEYWORDS = ('dall-e', 'tts', 'whisper', 'audio', 'moderation', 'davinci', 'babbage')

    async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
        """
        通过 /v1/models API 在线拉取最新模型列表
        
        适用于所有 OpenAI 兼容的提供商
        """
        import httpx

        base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
        headers = {}
        if self.api_key:
            headers['Authorization'] = f'Bearer {self.api_key}'

        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(f'{base_url}/models', headers=headers)
                resp.raise_for_status()
                data = resp.json()

            model_list = data.get('data', [])
            if not model_list:
                return []

            results = []
            for item in model_list:
                model_id = item.get('id', '')
                if not model_id:
                    continue

                model_lower = model_id.lower()

                # 跳过非文本模型
                if any(kw in model_lower for kw in self._SKIP_KEYWORDS):
                    continue

                # 推断模型类型
                if any(kw in model_lower for kw in self._RERANK_KEYWORDS):
                    model_type = 'rerank'
                elif any(kw in model_lower for kw in self._EMBEDDING_KEYWORDS):
                    model_type = 'embedding'
                else:
                    model_type = 'chat'

                results.append({
                    'model_name': model_id,
                    'display_name': model_id,
                    'model_type': model_type,
                    'max_tokens': 4096,
                    'context_window': 4096,
                    'supports_vision': False,
                    'supports_function_call': model_type == 'chat',
                    'input_price': 0,
                    'output_price': 0,
                })

            # 补充默认列表中的 embedding/rerank 模型（API 可能不返回这些类型）
            fetched_names = {m['model_name'] for m in results}
            for dm in self.__class__.get_default_models():
                if dm['model_type'] in ('embedding', 'rerank') and dm['model_name'] not in fetched_names:
                    results.append(dm)

            # 按类型排序：chat → embedding → rerank
            type_order = {'chat': 0, 'embedding': 1, 'rerank': 2}
            results.sort(key=lambda m: (type_order.get(m['model_type'], 9), m['model_name']))
            logger.info(f'从 {base_url} 拉取到 {len(results)} 个模型')
            return results

        except Exception as e:
            logger.warning(f'在线拉取模型列表失败 ({base_url}): {e}')
            return []

    @classmethod
    def get_default_models(cls) -> List[Dict[str, Any]]:
        """获取默认模型列表"""
        return [
            # ---- Chat 模型 ----
            {
                'model_name': 'o4-mini',
                'display_name': 'o4-mini',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.0011,
                'output_price': 0.0044,
            },
            {
                'model_name': 'o3-mini',
                'display_name': 'o3-mini',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 200000,
                'supports_vision': False,
                'supports_function_call': True,
                'input_price': 0.0011,
                'output_price': 0.0044,
            },
            {
                'model_name': 'gpt-4.1',
                'display_name': 'GPT-4.1',
                'model_type': 'chat',
                'max_tokens': 32768,
                'context_window': 1047576,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.002,
                'output_price': 0.008,
            },
            {
                'model_name': 'gpt-4.1-mini',
                'display_name': 'GPT-4.1 Mini',
                'model_type': 'chat',
                'max_tokens': 32768,
                'context_window': 1047576,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.0004,
                'output_price': 0.0016,
            },
            {
                'model_name': 'gpt-4.1-nano',
                'display_name': 'GPT-4.1 Nano',
                'model_type': 'chat',
                'max_tokens': 32768,
                'context_window': 1047576,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.0001,
                'output_price': 0.0004,
            },
            {
                'model_name': 'gpt-4o',
                'display_name': 'GPT-4o',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 128000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.0025,
                'output_price': 0.01,
            },
            {
                'model_name': 'gpt-4o-mini',
                'display_name': 'GPT-4o Mini',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 128000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.00015,
                'output_price': 0.0006,
            },
            # ---- Embedding 模型 ----
            {
                'model_name': 'text-embedding-3-small',
                'display_name': 'Text Embedding 3 Small',
                'model_type': 'embedding',
                'max_tokens': 8191,
                'context_window': 8191,
                'input_price': 0.00002,
                'output_price': 0,
            },
            {
                'model_name': 'text-embedding-3-large',
                'display_name': 'Text Embedding 3 Large',
                'model_type': 'embedding',
                'max_tokens': 8191,
                'context_window': 8191,
                'input_price': 0.00013,
                'output_price': 0,
            },
            {
                'model_name': 'text-embedding-ada-002',
                'display_name': 'Text Embedding Ada 002',
                'model_type': 'embedding',
                'max_tokens': 8191,
                'context_window': 8191,
                'input_price': 0.0001,
                'output_price': 0,
            },
        ]
