"""
AI瀹㈡埛绔?- 鏁村悎鎵€鏈堿I鍔熻兘
"""
import inspect
import json
import re
from typing import List, Optional, Dict, Any, AsyncIterator, Tuple
from pathlib import Path
from .base import ModelConfig, ModelProvider, Message, ToolRegistry
from .models import OpenAIModel, AnthropicModel
from .personality import PersonalitySystem
from .memory import MemorySystem
from .task_manager import LongTaskManager
from src.utils.logger import setup_logger

logger = setup_logger('AIClient')


class AIClient:
    """AI瀹㈡埛绔?- 缁熶竴鎺ュ彛"""
    
    def __init__(
        self,
        model_config: ModelConfig,
        embed_config: Optional[ModelConfig] = None,
        data_dir: Path = Path("data/ai"),
        use_vector_db: bool = True
    ):
        self.config = model_config
        self.data_dir = data_dir
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
        # 初始化主模型
        self.model = self._create_model(model_config)
        
        # 初始化嵌入模型（如果提供）
        self.embed_model = None
        if embed_config:
            self.embed_model = self._create_model(embed_config)
            logger.info(
                f"嵌入模型初始化完成: {embed_config.provider.value}/{embed_config.model_name}"
            )
        
        # 初始化工具注册表
        self.tools = ToolRegistry()
        self._tool_sources: Dict[str, str] = {}
        
        # 初始化人格系统
        self.personality = PersonalitySystem(
            config_path=data_dir / "personalities.json"
        )
        
        # 初始化记忆系统
        self.memory = MemorySystem(
            storage_path=data_dir / "long_term_memory.json",
            embed_func=self._embed_wrapper,
            importance_evaluator=self._evaluate_memory_importance,
            use_vector_db=use_vector_db
        )
        
        # 初始化长任务管理器
        self.task_manager = LongTaskManager(
            storage_path=data_dir / "tasks.json"
        )
        
        logger.info(
            f"AI 客户端初始化完成: {model_config.provider.value}/{model_config.model_name}"
        )
    
    def _create_model(self, config: ModelConfig):
        """创建模型实例。"""
        if config.provider == ModelProvider.OPENAI:
            return OpenAIModel(config)
        elif config.provider == ModelProvider.ANTHROPIC:
            return AnthropicModel(config)
        elif config.provider in [ModelProvider.DEEPSEEK, ModelProvider.QWEN]:
            # DeepSeek 和 Qwen 使用 OpenAI 兼容接口
            return OpenAIModel(config)
        else:
            raise ValueError(f"不支持的模型提供商: {config.provider}")
    
    async def _embed_wrapper(self, text: str) -> List[float]:
        """嵌入向量包装器。"""
        try:
            # 如果有独立的嵌入模型，优先使用
            if self.embed_model:
                return await self.embed_model.embed(text)
            # 否则尝试使用主模型
            return await self.model.embed(text)
        except NotImplementedError:
            # 如果都不支持嵌入，返回 None（记忆系统会降级）
            logger.warning("Current model does not support embeddings; vector retrieval disabled")
            return None
        except Exception as e:
            logger.error(f"生成嵌入向量失败: {e}")
            return None

    @staticmethod
    def _parse_importance_score(raw: str) -> float:
        text = (raw or "").strip()
        if not text:
            raise ValueError("empty importance response")

        try:
            parsed = json.loads(text)
            if isinstance(parsed, (int, float)):
                return float(parsed)
            if isinstance(parsed, dict):
                for key in ["importance", "score", "value"]:
                    if key in parsed:
                        return float(parsed[key])
        except Exception:
            pass

        match = re.search(r"-?\d+(?:\.\d+)?", text)
        if not match:
            raise ValueError(f"cannot parse importance score: {text}")
        return float(match.group(0))

    async def _evaluate_memory_importance(
        self, content: str, metadata: Optional[Dict] = None
    ) -> float:
        """
        调用主模型评估记忆重要性，返回 [0, 1] 分值。
        """
        system_prompt = (
            "你是记忆重要性评估器。请根据输入内容判断该信息是否值得长期记忆。"
            "输出一个 0 到 1 的数字，数字越大表示越重要。"
            "只输出数字，不要输出任何解释、单位或多余文本。"
        )
        payload = json.dumps(
            {"content": content, "metadata": metadata or {}},
            ensure_ascii=False,
        )
        messages = [
            Message(role="system", content=system_prompt),
            Message(role="user", content=payload),
        ]

        try:
            response = await self.model.chat(
                messages=messages,
                tools=None,
                temperature=0.0,
                max_tokens=16,
            )
            score = self._parse_importance_score(response.content)
            return max(0.0, min(1.0, score))
        except Exception as e:
            logger.warning(f"memory importance evaluation failed, fallback to neutral score: {e}")
            return 0.5
    
    async def chat(
        self,
        user_id: str,
        user_message: str,
        system_prompt: Optional[str] = None,
        use_memory: bool = True,
        use_tools: bool = True,
        stream: bool = False,
        **kwargs
    ) -> str:
        """对话接口。"""
        try:
            # 构建消息列表
            messages = []
            
            # 系统提示词
            if system_prompt is None:
                system_prompt = self.personality.get_system_prompt()
            
            # 注入记忆上下文
            if use_memory:
                short_term, long_term = await self.memory.get_context(
                    user_id=user_id,
                    query=user_message
                )
                
                if short_term or long_term:
                    memory_context = self.memory.format_context(short_term, long_term)
                    system_prompt += f"\n\n{memory_context}"
            
            messages.append(Message(role="system", content=system_prompt))
            
            # 添加用户消息
            messages.append(Message(role="user", content=user_message))
            
            # 准备工具
            tools = None
            if use_tools and self.tools.list():
                tools = self.tools.to_openai_format()
            
            # 调用模型
            if stream:
                return self._chat_stream(messages, tools, **kwargs)
            else:
                response = await self.model.chat(messages, tools, **kwargs)
                
                # 处理工具调用
                if response.tool_calls:
                    response = await self._handle_tool_calls(
                        messages, response, tools, **kwargs
                    )
                
                # 写入记忆
                if use_memory:
                    stored_memory = await self.memory.add_qa_pair(
                        user_id=user_id,
                        question=user_message,
                        answer=response.content,
                        metadata={"source": "chat"},
                    )
                    if stored_memory:
                        logger.info(
                            "已写入长期记忆问答对:\n"
                            f"{stored_memory.content}\n"
                            f"memory_id={stored_memory.id}, "
                            f"importance={stored_memory.importance:.2f}"
                        )
                
                return response.content
                
        except Exception as e:
            logger.error(f"对话失败: {e}")
            raise
    
    async def _chat_stream(
        self,
        messages: List[Message],
        tools: Optional[List[Dict]],
        **kwargs
    ) -> AsyncIterator[str]:
        """流式对话。"""
        async for chunk in self.model.chat_stream(messages, tools, **kwargs):
            yield chunk
    
    async def _handle_tool_calls(
        self,
        messages: List[Message],
        response: Message,
        tools: Optional[List[Dict]],
        **kwargs
    ) -> Message:
        """处理工具调用。"""
        messages.append(response)
        total_calls = len(response.tool_calls or [])
        if total_calls:
            logger.info(f"检测到工具调用请求: {total_calls} 个")
        
        # 执行工具调用
        for tool_call in response.tool_calls or []:
            try:
                tool_name, tool_args, tool_call_id = self._parse_tool_call(tool_call)
            except Exception as e:
                logger.warning(f"解析工具调用失败: {e}")
                fallback_id = tool_call.get('id') if isinstance(tool_call, dict) else getattr(tool_call, 'id', None)
                if fallback_id:
                    messages.append(Message(
                        role="tool",
                        content=f"工具参数解析失败: {str(e)}",
                        tool_call_id=fallback_id,
                        name="tool"
                    ))
                continue
            if not tool_name:
                logger.warning(f"跳过无效工具调用: {tool_call}")
                continue
            
            tool_def = self.tools.get(tool_name)
            tool_source = self._tool_sources.get(tool_name, "custom")
            if not tool_def:
                error_msg = f"未找到工具: {tool_name}"
                logger.warning(error_msg)
                messages.append(Message(
                    role="tool",
                    name=tool_name,
                    content=error_msg,
                    tool_call_id=tool_call_id
                ))
                continue
            
            try:
                logger.info(
                    "工具调用开始: "
                    f"source={tool_source}, name={tool_name}, "
                    f"args={self._preview_log_payload(tool_args)}"
                )
                result = tool_def.function(**tool_args)
                if inspect.isawaitable(result):
                    result = await result
                logger.info(
                    "工具调用成功: "
                    f"source={tool_source}, name={tool_name}, "
                    f"result={self._preview_log_payload(result)}"
                )
                messages.append(Message(
                    role="tool",
                    name=tool_name,
                    content=str(result),
                    tool_call_id=tool_call_id
                ))
            except Exception as e:
                logger.warning(
                    "工具调用失败: "
                    f"source={tool_source}, name={tool_name}, error={e}"
                )
                messages.append(Message(
                    role="tool",
                    name=tool_name,
                    content=f"工具执行失败: {str(e)}",
                    tool_call_id=tool_call_id
                ))

        # 再次调用模型获取最终响应
        return await self.model.chat(messages, tools, **kwargs)

    def _parse_tool_call(self, tool_call: Any) -> Tuple[Optional[str], Dict[str, Any], Optional[str]]:
        """兼容不同 SDK 返回的工具调用结构。"""
        if isinstance(tool_call, dict):
            tool_call_id = tool_call.get('id')
            function = tool_call.get('function') or {}
            tool_name = function.get('name')
            raw_args = function.get('arguments')
        else:
            tool_call_id = getattr(tool_call, 'id', None)
            function = getattr(tool_call, 'function', None)
            tool_name = getattr(function, 'name', None) if function else None
            raw_args = getattr(function, 'arguments', None) if function else None
        
        tool_args = self._normalize_tool_args(raw_args)
        return tool_name, tool_args, tool_call_id

    def _normalize_tool_args(self, raw_args: Any) -> Dict[str, Any]:
        """将工具参数统一转换为字典。"""
        if raw_args is None:
            return {}
        
        if isinstance(raw_args, dict):
            return raw_args
        
        if isinstance(raw_args, str):
            raw_args = raw_args.strip()
            if not raw_args:
                return {}
            parsed = json.loads(raw_args)
            if not isinstance(parsed, dict):
                raise ValueError(f"工具参数必须是 JSON 对象，实际类型: {type(parsed)}")
            return parsed
        
        if hasattr(raw_args, 'model_dump'):
            parsed = raw_args.model_dump()
            if isinstance(parsed, dict):
                return parsed
        
        raise ValueError(f"不支持的工具参数类型: {type(raw_args)}")

    @staticmethod
    def _preview_log_payload(payload: Any, max_len: int = 240) -> str:
        """日志中展示参数/结果时使用的简短预览。"""
        try:
            text = json.dumps(payload, ensure_ascii=False, default=str)
        except Exception:
            text = str(payload)

        if len(text) > max_len:
            return text[:max_len] + "..."
        return text
    
    def set_personality(self, personality_name: str) -> bool:
        """设置人格。"""
        return self.personality.set_personality(personality_name)
    
    def list_personalities(self) -> List[str]:
        """列出所有人格。"""
        return self.personality.list_personalities()

    def switch_model(self, model_config: ModelConfig) -> bool:
        """Runtime switch for primary chat model."""
        new_model = self._create_model(model_config)
        self.model = new_model
        self.config = model_config
        logger.info(
            f"已切换主模型: {model_config.provider.value}/{model_config.model_name}"
        )
        return True

    async def create_long_task(
        self,
        user_id: str,
        title: str,
        description: str,
        steps: List[Dict],
        metadata: Optional[Dict] = None
    ) -> str:
        """创建长任务。"""
        return self.task_manager.create_task(
            user_id=user_id,
            title=title,
            description=description,
            steps=steps,
            metadata=metadata
        )
    
    async def start_task(
        self,
        task_id: str,
        progress_callback: Optional[callable] = None
    ):
        """启动任务。"""
        await self.task_manager.start_task(task_id, progress_callback)
    
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """获取任务状态。"""
        return self.task_manager.get_task_status(task_id)
    
    def register_tool(
        self,
        name: str,
        description: str,
        parameters: Dict,
        function: callable,
        source: str = "custom",
    ):
        """注册工具。"""
        from .base import ToolDefinition
        tool = ToolDefinition(
            name=name,
            description=description,
            parameters=parameters,
            function=function
        )
        self.tools.register(tool)
        self._tool_sources[name] = source
        logger.info(f"已注册工具: {name} (source={source})")

    def unregister_tool(self, name: str) -> bool:
        """卸载工具。"""
        removed = self.tools.unregister(name)
        if removed:
            self._tool_sources.pop(name, None)
            logger.info(f"已卸载工具: {name}")
        return removed

    def unregister_tools_by_prefix(self, prefix: str) -> int:
        """按前缀批量卸载工具。"""
        removed_count = self.tools.unregister_by_prefix(prefix)
        for tool_name in list(self._tool_sources.keys()):
            if tool_name.startswith(prefix):
                self._tool_sources.pop(tool_name, None)
        if removed_count:
            logger.info(f"Unregistered tools by prefix {prefix}: {removed_count}")
        return removed_count
    
    def clear_memory(self, user_id: str):
        """清除用户短期记忆。"""
        self.memory.clear_short_term(user_id)
        logger.info(f"Cleared short-term memory for user {user_id}")

    async def clear_long_term_memory(self, user_id: str) -> bool:
        try:
            await self.memory.clear_long_term(user_id)
            logger.info(f"Cleared long-term memory for user {user_id}")
            return True
        except Exception as e:
            logger.warning(f"Failed to clear long-term memory for user {user_id}: {e}")
            return False

    async def list_long_term_memories(self, user_id: str, limit: int = 20):
        return await self.memory.list_long_term(user_id, limit=limit)

    async def get_long_term_memory(self, user_id: str, memory_id: str):
        return await self.memory.get_long_term(user_id, memory_id)

    async def add_long_term_memory(
        self,
        user_id: str,
        content: str,
        importance: float = 0.8,
        metadata: Optional[Dict] = None,
    ):
        return await self.memory.add_long_term(
            user_id=user_id,
            content=content,
            importance=importance,
            metadata=metadata,
        )

    async def search_long_term_memories(
        self, user_id: str, query: str, limit: int = 10
    ):
        return await self.memory.search_long_term(user_id, query=query, limit=limit)

    async def update_long_term_memory(
        self,
        user_id: str,
        memory_id: str,
        content: Optional[str] = None,
        importance: Optional[float] = None,
        metadata: Optional[Dict] = None,
    ):
        return await self.memory.update_long_term(
            user_id=user_id,
            memory_id=memory_id,
            content=content,
            importance=importance,
            metadata=metadata,
        )

    async def delete_long_term_memory(self, user_id: str, memory_id: str) -> bool:
        return await self.memory.delete_long_term(user_id, memory_id)

    async def clear_all_memory(self, user_id: str) -> bool:
        """清除用户全部记忆（短期 + 长期）。"""
        self.clear_memory(user_id)
        try:
            return await self.clear_long_term_memory(user_id)
        except Exception:
            return False

