"""
Anthropic Claude 提供商适配器
"""
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 ClaudeProvider(BaseLLMProvider):
    """
    Anthropic Claude 提供商适配器
    
    支持 Claude 3 系列模型
    """
    
    provider_type = 'claude'
    provider_name = 'Anthropic Claude'
    supported_model_types = ['chat']
    
    DEFAULT_API_BASE = 'https://api.anthropic.com'
    
    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:
                import anthropic
                self._client = anthropic.Anthropic(
                    api_key=self.api_key,
                    base_url=self.api_base if self.api_base != self.DEFAULT_API_BASE else None,
                )
            except ImportError:
                raise ImportError('请安装 anthropic 库: pip install anthropic')
        return self._client
    
    def _get_async_client(self):
        """获取异步客户端"""
        if self._async_client is None:
            try:
                import anthropic
                self._async_client = anthropic.AsyncAnthropic(
                    api_key=self.api_key,
                    base_url=self.api_base if self.api_base != self.DEFAULT_API_BASE else None,
                )
            except ImportError:
                raise ImportError('请安装 anthropic 库: pip install anthropic')
        return self._async_client
    
    def _convert_messages(self, messages: List[LLMMessage]) -> tuple:
        """
        转换消息格式
        
        Claude API 需要将 system 消息单独传递
        """
        system_prompt = ''
        converted_messages = []
        
        for msg in messages:
            if msg.role == 'system':
                system_prompt = msg.content
            elif msg.role == 'tool':
                # Claude 的工具结果消息格式
                converted_messages.append({
                    'role': 'user',
                    'content': [{
                        'type': 'tool_result',
                        'tool_use_id': msg.tool_call_id,
                        'content': msg.content,
                    }],
                })
            elif msg.role == 'assistant' and msg.tool_calls:
                # 包含工具调用的助手消息
                content = []
                if msg.content:
                    content.append({'type': 'text', 'text': msg.content})
                for tc in msg.tool_calls:
                    content.append({
                        'type': 'tool_use',
                        'id': tc.get('id', ''),
                        'name': tc.get('function', {}).get('name', ''),
                        'input': tc.get('function', {}).get('arguments', {}),
                    })
                converted_messages.append({
                    'role': 'assistant',
                    'content': content,
                })
            else:
                converted_messages.append({
                    'role': msg.role,
                    'content': msg.content,
                })
        
        return system_prompt, converted_messages
    
    def chat(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> LLMResponse:
        """同步对话"""
        client = self._get_client()
        system_prompt, converted_messages = self._convert_messages(messages)
        
        kwargs = {
            'model': config.model,
            'messages': converted_messages,
            'max_tokens': config.max_tokens,
            'temperature': config.temperature,
            'top_p': config.top_p,
        }
        
        if system_prompt:
            kwargs['system'] = system_prompt
        
        if config.stop:
            kwargs['stop_sequences'] = config.stop
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_claude_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = {'type': 'any'}
            elif config.tool_choice == 'none':
                # Claude 不支持 none，不传递 tools 即可
                del kwargs['tools']
            elif config.tool_choice != 'auto':
                # 指定具体工具
                kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
            else:
                kwargs['tool_choice'] = {'type': 'auto'}
        
        response = client.messages.create(**kwargs)
        
        # 解析响应内容和工具调用
        content = ''
        tool_calls = []
        
        for block in response.content:
            if block.type == 'text':
                content = block.text
            elif block.type == 'tool_use':
                tool_calls.append(ToolCall(
                    id=block.id,
                    name=block.name,
                    arguments=block.input if isinstance(block.input, dict) else {},
                ))
        
        return LLMResponse(
            content=content,
            model=response.model,
            prompt_tokens=response.usage.input_tokens,
            completion_tokens=response.usage.output_tokens,
            total_tokens=response.usage.input_tokens + response.usage.output_tokens,
            finish_reason=response.stop_reason or '',
            raw_response=response.model_dump(),
            tool_calls=tool_calls if tool_calls else None,
        )
    
    async def chat_async(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> LLMResponse:
        """异步对话"""
        client = self._get_async_client()
        system_prompt, converted_messages = self._convert_messages(messages)
        
        kwargs = {
            'model': config.model,
            'messages': converted_messages,
            'max_tokens': config.max_tokens,
            'temperature': config.temperature,
            'top_p': config.top_p,
        }
        
        if system_prompt:
            kwargs['system'] = system_prompt
        
        if config.stop:
            kwargs['stop_sequences'] = config.stop
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_claude_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = {'type': 'any'}
            elif config.tool_choice == 'none':
                del kwargs['tools']
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
            else:
                kwargs['tool_choice'] = {'type': 'auto'}
        
        response = await client.messages.create(**kwargs)
        
        # 解析响应内容和工具调用
        content = ''
        tool_calls = []
        
        for block in response.content:
            if block.type == 'text':
                content = block.text
            elif block.type == 'tool_use':
                tool_calls.append(ToolCall(
                    id=block.id,
                    name=block.name,
                    arguments=block.input if isinstance(block.input, dict) else {},
                ))
        
        return LLMResponse(
            content=content,
            model=response.model,
            prompt_tokens=response.usage.input_tokens,
            completion_tokens=response.usage.output_tokens,
            total_tokens=response.usage.input_tokens + response.usage.output_tokens,
            finish_reason=response.stop_reason or '',
            raw_response=response.model_dump(),
            tool_calls=tool_calls if tool_calls else None,
        )
    
    async def chat_stream(
        self,
        messages: List[LLMMessage],
        config: LLMConfig,
    ) -> AsyncGenerator[LLMStreamChunk, None]:
        """流式对话"""
        client = self._get_async_client()
        system_prompt, converted_messages = self._convert_messages(messages)
        
        kwargs = {
            'model': config.model,
            'messages': converted_messages,
            'max_tokens': config.max_tokens,
            'temperature': config.temperature,
            'top_p': config.top_p,
        }
        
        if system_prompt:
            kwargs['system'] = system_prompt
        
        if config.stop:
            kwargs['stop_sequences'] = config.stop
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_claude_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = {'type': 'any'}
            elif config.tool_choice == 'none':
                del kwargs['tools']
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
            else:
                kwargs['tool_choice'] = {'type': 'auto'}
        
        # 用于累积工具调用
        tool_calls_accumulator = []  # List of {id, name, input_json}
        current_tool_call = None
        
        async with client.messages.stream(**kwargs) as stream:
            async for event in stream:
                if hasattr(event, 'type'):
                    if event.type == 'content_block_start':
                        if hasattr(event, 'content_block') and event.content_block.type == 'tool_use':
                            current_tool_call = {
                                'id': event.content_block.id,
                                'name': event.content_block.name,
                                'input_json': '',
                            }
                    elif event.type == 'content_block_delta':
                        if hasattr(event, 'delta'):
                            if event.delta.type == 'text_delta':
                                yield LLMStreamChunk(
                                    content=event.delta.text,
                                    is_finished=False,
                                )
                            elif event.delta.type == 'input_json_delta' and current_tool_call:
                                current_tool_call['input_json'] += event.delta.partial_json
                    elif event.type == 'content_block_stop':
                        if current_tool_call:
                            tool_calls_accumulator.append(current_tool_call)
                            current_tool_call = None
            
            # 获取最终的 usage 信息
            final_message = await stream.get_final_message()
            
            # 解析工具调用
            tool_calls = None
            if tool_calls_accumulator:
                tool_calls = self._parse_accumulated_tool_calls(tool_calls_accumulator)
            
            yield LLMStreamChunk(
                content='',
                is_finished=True,
                finish_reason=final_message.stop_reason or '',
                prompt_tokens=final_message.usage.input_tokens,
                completion_tokens=final_message.usage.output_tokens,
                total_tokens=final_message.usage.input_tokens + final_message.usage.output_tokens,
                tool_calls=tool_calls,
            )
    
    def _parse_accumulated_tool_calls(self, accumulator: List[Dict]) -> List[ToolCall]:
        """解析累积的工具调用"""
        import json
        result = []
        for tc in accumulator:
            try:
                arguments = json.loads(tc['input_json']) if tc['input_json'] else {}
            except json.JSONDecodeError:
                arguments = {'raw': tc['input_json']}
            
            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()
        system_prompt, converted_messages = self._convert_messages(messages)
        
        kwargs = {
            'model': config.model,
            'messages': converted_messages,
            'max_tokens': config.max_tokens,
            'temperature': config.temperature,
            'top_p': config.top_p,
        }
        
        if system_prompt:
            kwargs['system'] = system_prompt
        
        if config.stop:
            kwargs['stop_sequences'] = config.stop
        
        # 添加 Function Calling 参数
        if config.tools:
            kwargs['tools'] = [t.to_claude_format() for t in config.tools]
            if config.tool_choice == 'required':
                kwargs['tool_choice'] = {'type': 'any'}
            elif config.tool_choice == 'none':
                del kwargs['tools']
            elif config.tool_choice != 'auto':
                kwargs['tool_choice'] = {'type': 'tool', 'name': config.tool_choice}
            else:
                kwargs['tool_choice'] = {'type': 'auto'}
        
        # 用于累积工具调用
        tool_calls_accumulator = []
        current_tool_call = None
        
        with client.messages.stream(**kwargs) as stream:
            for event in stream:
                if hasattr(event, 'type'):
                    if event.type == 'content_block_start':
                        if hasattr(event, 'content_block') and event.content_block.type == 'tool_use':
                            current_tool_call = {
                                'id': event.content_block.id,
                                'name': event.content_block.name,
                                'input_json': '',
                            }
                    elif event.type == 'content_block_delta':
                        if hasattr(event, 'delta'):
                            if event.delta.type == 'text_delta':
                                yield LLMStreamChunk(
                                    content=event.delta.text,
                                    is_finished=False,
                                )
                            elif event.delta.type == 'input_json_delta' and current_tool_call:
                                current_tool_call['input_json'] += event.delta.partial_json
                    elif event.type == 'content_block_stop':
                        if current_tool_call:
                            tool_calls_accumulator.append(current_tool_call)
                            current_tool_call = None
            
            # 获取最终的 usage 信息
            final_message = stream.get_final_message()
            
            # 解析工具调用
            tool_calls = None
            if tool_calls_accumulator:
                tool_calls = self._parse_accumulated_tool_calls(tool_calls_accumulator)
            
            yield LLMStreamChunk(
                content='',
                is_finished=True,
                finish_reason=final_message.stop_reason or '',
                prompt_tokens=final_message.usage.input_tokens,
                completion_tokens=final_message.usage.output_tokens,
                total_tokens=final_message.usage.input_tokens + final_message.usage.output_tokens,
                tool_calls=tool_calls,
            )
    
    async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
        """通过 Anthropic /v1/models 端点在线拉取模型列表"""
        import httpx

        base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
        headers = {
            'x-api-key': self.api_key,
            'anthropic-version': '2023-06-01',
        }

        try:
            async with httpx.AsyncClient(timeout=15) as client:
                resp = await client.get(f'{base_url}/v1/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', '')
                display = item.get('display_name', model_id)
                if not model_id:
                    continue

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

            results.sort(key=lambda m: m['model_name'])
            logger.info(f'从 Anthropic ({base_url}) 拉取到 {len(results)} 个模型')
            return results

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

    @classmethod
    def get_default_models(cls) -> List[Dict[str, Any]]:
        """获取默认模型列表"""
        return [
            # ---- Chat 模型 ----
            {
                'model_name': 'claude-sonnet-4-20250514',
                'display_name': 'Claude Sonnet 4',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.003,
                'output_price': 0.015,
            },
            {
                'model_name': 'claude-3-7-sonnet-20250219',
                'display_name': 'Claude 3.7 Sonnet',
                'model_type': 'chat',
                'max_tokens': 16384,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.003,
                'output_price': 0.015,
            },
            {
                'model_name': 'claude-3-5-sonnet-20241022',
                'display_name': 'Claude 3.5 Sonnet',
                'model_type': 'chat',
                'max_tokens': 8192,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.003,
                'output_price': 0.015,
            },
            {
                'model_name': 'claude-3-5-haiku-20241022',
                'display_name': 'Claude 3.5 Haiku',
                'model_type': 'chat',
                'max_tokens': 8192,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.001,
                'output_price': 0.005,
            },
            {
                'model_name': 'claude-3-opus-20240229',
                'display_name': 'Claude 3 Opus',
                'model_type': 'chat',
                'max_tokens': 4096,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.015,
                'output_price': 0.075,
            },
            {
                'model_name': 'claude-3-haiku-20240307',
                'display_name': 'Claude 3 Haiku',
                'model_type': 'chat',
                'max_tokens': 4096,
                'context_window': 200000,
                'supports_vision': True,
                'supports_function_call': True,
                'input_price': 0.00025,
                'output_price': 0.00125,
            },
        ]
