class PerceptionModule:
    """Agent 感知模块"""
    
    def __init__(self, llm_client):
        self.llm = llm_client
        self.parsers = {
            "text": TextParser(),
            "image": ImageParser(),
            "audio": AudioParser(),
            "structured": StructuredDataParser()
        }
    
    def perceive(self, input_data: dict) -> str:
        """将多模态输入统一转化为文本描述"""
        input_type = input_data.get("type", "text")
        parser = self.parsers.get(input_type)
        
        if parser is None:
            raise ValueError(f"不支持的输入类型: {input_type}")
        
        # 解析输入
        parsed = parser.parse(input_data["content"])
        
        # 提取关键信息
        context = self._extract_context(parsed)
        
        # 意图识别
        intent = self._identify_intent(context)
        
        return {
            "raw_input": input_data,
            "parsed_content": parsed,
            "context": context,
            "intent": intent
        }
    
    def _identify_intent(self, context: str) -> dict:
        """识别用户意图"""
        prompt = f"""分析以下输入的用户意图：
        
输入内容：{context}

请返回：
1. 主要意图（一句话描述）
2. 任务类型（查询/创建/修改/删除/分析/其他）
3. 关键实体（涉及的对象）
4. 约束条件（时间、范围等限制）"""
        
        return self.llm.analyze(prompt)

class PlanningModule:
    """Agent 规划模块"""
    
    def __init__(self, llm_client, tool_registry):
        self.llm = llm_client
        self.tools = tool_registry
    
    def create_plan(self, task: str, context: dict) -> list:
        """创建执行计划"""
        available_tools = self.tools.get_tool_descriptions()
        
        prompt = f"""你是一个任务规划专家。请为以下任务创建详细的执行计划。

任务描述：{task}

上下文信息：{context}

可用工具：
{available_tools}

请按以下格式输出执行计划：
1. [步骤描述] - 使用工具：[工具名] - 依赖步骤：[无/步骤编号]
2. ...

要求：
- 步骤要具体、可执行
- 明确每步使用的工具
- 标注步骤间的依赖关系
- 考虑可能的失败情况和备选方案"""
        
        plan = self.llm.generate(prompt)
        return self._parse_plan(plan)
    
    def replan(self, original_plan: list, execution_result: dict, 
               error_info: str) -> list:
        """根据执行结果动态重规划"""
        prompt = f"""原始计划执行遇到问题，需要重新规划。

原始计划：{original_plan}
已执行结果：{execution_result}
错误信息：{error_info}

请分析失败原因，并生成新的执行计划。
保留已成功的步骤，只调整失败和后续步骤。"""
        
        new_plan = self.llm.generate(prompt)
        return self._parse_plan(new_plan)

class MemoryModule:
    """Agent 记忆模块"""
    
    def __init__(self, vector_store, embedding_model):
        self.short_term = []          # 短期记忆（对话历史）
        self.working = {}             # 工作记忆（任务状态）
        self.vector_store = vector_store  # 长期记忆（向量数据库）
        self.embedding = embedding_model
    
    # 短期记忆操作
    def add_to_short_term(self, role: str, content: str):
        """添加到短期记忆"""
        self.short_term.append({"role": role, "content": content})
        # 超过上下文窗口限制时，压缩旧记忆
        if self._estimate_tokens() > 100000:
            self._compress_short_term()
    
    def _compress_short_term(self):
        """压缩短期记忆，保留关键信息"""
        # 保留最近 10 轮对话
        recent = self.short_term[-20:]
        # 将旧对话总结后存入长期记忆
        old = self.short_term[:-20]
        summary = self._summarize(old)
        self.save_to_long_term(summary, category="conversation_summary")
        self.short_term = [{"role": "system", "content": f"历史对话摘要：{summary}"}] + recent
    
    # 长期记忆操作
    def save_to_long_term(self, content: str, category: str = "general",
                          metadata: dict = None):
        """保存到长期记忆"""
        embedding = self.embedding.encode(content)
        self.vector_store.upsert({
            "content": content,
            "embedding": embedding,
            "category": category,
            "metadata": metadata or {},
            "timestamp": datetime.now().isoformat()
        })
    
    def recall(self, query: str, top_k: int = 5, 
               category: str = None) -> list:
        """从长期记忆中检索相关信息"""
        query_embedding = self.embedding.encode(query)
        filters = {"category": category} if category else {}
        results = self.vector_store.search(
            query_embedding, top_k=top_k, filters=filters
        )
        return results
    
    # 工作记忆操作
    def update_working(self, key: str, value):
        """更新工作记忆"""
        self.working[key] = value
    
    def get_working(self, key: str, default=None):
        """获取工作记忆"""
        return self.working.get(key, default)

class ActionModule:
    """Agent 行动模块"""
    
    def __init__(self, tool_registry):
        self.tools = tool_registry
        self.execution_history = []
    
    def execute(self, action: dict) -> dict:
        """执行一个行动"""
        action_type = action["type"]
        
        try:
            if action_type == "tool_call":
                result = self._execute_tool(action)
            elif action_type == "code_execution":
                result = self._execute_code(action)
            elif action_type == "api_request":
                result = self._execute_api(action)
            elif action_type == "human_interaction":
                result = self._ask_human(action)
            else:
                raise ValueError(f"未知的行动类型: {action_type}")
            
            # 记录执行历史
            self.execution_history.append({
                "action": action,
                "result": result,
                "status": "success",
                "timestamp": datetime.now().isoformat()
            })
            
            return {"status": "success", "result": result}
            
        except Exception as e:
            error_result = {
                "status": "error",
                "error": str(e),
                "action": action
            }
            self.execution_history.append(error_result)
            return error_result
    
    def _execute_tool(self, action: dict) -> any:
        """执行工具调用"""
        tool_name = action["tool_name"]
        tool_args = action["arguments"]
        
        tool = self.tools.get(tool_name)
        if tool is None:
            raise ValueError(f"工具不存在: {tool_name}")
        
        # 参数校验
        tool.validate_args(tool_args)
        
        # 执行工具
        return tool.execute(**tool_args)

class ToolRegistry:
    """工具注册中心"""
    
    def __init__(self):
        self.tools = {}
    
    def register(self, name: str, func: callable, description: str,
                 parameters: dict):
        """注册一个工具"""
        self.tools[name] = {
            "function": func,
            "description": description,
            "parameters": parameters
        }
    
    def get_tool_descriptions(self) -> str:
        """获取所有工具的描述（用于 Prompt）"""
        descriptions = []
        for name, tool in self.tools.items():
            desc = f"- {name}: {tool['description']}\n"
            desc += f"  参数: {json.dumps(tool['parameters'], ensure_ascii=False)}"
            descriptions.append(desc)
        return "\n".join(descriptions)
    
    def execute(self, name: str, **kwargs) -> any:
        """执行指定工具"""
        tool = self.tools.get(name)
        if not tool:
            raise ValueError(f"工具 {name} 未注册")
        return tool["function"](**kwargs)


# 注册常用工具示例
registry = ToolRegistry()

# 搜索工具
registry.register(
    name="web_search",
    func=search_web,
    description="搜索互联网获取最新信息",
    parameters={
        "query": {"type": "string", "description": "搜索关键词"},
        "max_results": {"type": "integer", "description": "最大结果数", "default": 5}
    }
)

# 代码执行工具
registry.register(
    name="execute_python",
    func=run_python_code,
    description="执行 Python 代码并返回结果",
    parameters={
        "code": {"type": "string", "description": "要执行的 Python 代码"}
    }
)

# 数据库查询工具
registry.register(
    name="query_database",
    func=query_db,
    description="执行 SQL 查询",
    parameters={
        "sql": {"type": "string", "description": "SQL 查询语句"},
        "database": {"type": "string", "description": "数据库名称"}
    }
)

class LLMBrain:
    """Agent 的 LLM 大脑"""
    
    def __init__(self, model: str = "claude-3-5-sonnet", 
                 temperature: float = 0.0):
        self.model = model
        self.temperature = temperature
        self.client = get_llm_client(model)
    
    def think(self, system_prompt: str, messages: list, 
              tools: list = None) -> dict:
        """Agent 思考（LLM 推理）"""
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                *messages
            ],
            tools=tools,
            temperature=self.temperature,
            max_tokens=4096
        )
        return self._parse_response(response)
    
    def _parse_response(self, response) -> dict:
        """解析 LLM 响应"""
        choice = response.choices[0]
        
        if choice.finish_reason == "tool_calls":
            # LLM 决定调用工具
            return {
                "type": "tool_call",
                "tool_calls": [
                    {
                        "name": tc.function.name,
                        "arguments": json.loads(tc.function.arguments)
                    }
                    for tc in choice.message.tool_calls
                ]
            }
        else:
            # LLM 直接回复
            return {
                "type": "text",
                "content": choice.message.content
            }

AGENT_SYSTEM_PROMPT = """
# 角色定义
你是一个专业的数据分析 Agent，能够帮助用户完成数据查询、分析和可视化任务。

# 核心能力
1. 理解用户的数据分析需求
2. 编写 SQL 查询语句
3. 执行数据分析代码
4. 生成可视化图表
5. 撰写分析报告

# 可用工具
- query_database: 执行 SQL 查询，参数为 sql 语句和数据库名
- execute_python: 执行 Python 代码，用于数据处理和可视化
- web_search: 搜索互联网获取参考信息

# 工作流程
1. 理解用户需求，明确分析目标
2. 确定需要查询的数据表和字段
3. 编写并执行 SQL 查询
4. 对查询结果进行分析处理
5. 生成可视化图表（如需要）
6. 撰写分析结论和建议

# 行为准则
- 执行 SQL 前先确认表结构
- 大数据量查询要添加 LIMIT 限制
- 分析结论要有数据支撑
- 遇到错误要分析原因并重试
- 不确定的信息要向用户确认

# 输出格式
- 分析过程要清晰展示
- 关键数据要用表格呈现
- 结论要简洁明了
"""

# Zero-shot CoT
prompt = """
请帮我分析这个 SQL 查询为什么慢。

查询语句：
SELECT * FROM orders WHERE user_id = 123 AND create_time > '2025-01-01'

请一步步分析：
1. 首先检查查询条件...
2. 然后分析索引使用情况...
3. 最后给出优化建议...
"""

# Few-shot CoT（提供示例）
prompt = """
示例：
问题：为什么这个 API 响应慢？
思考过程：
1. 检查网络延迟 → 正常（50ms）
2. 检查数据库查询 → 发现全表扫描（3s）
3. 检查索引 → 缺少 user_id 索引
4. 结论：添加索引可以将查询时间从 3s 降到 50ms
建议：ALTER TABLE orders ADD INDEX idx_user_id (user_id);

现在请分析：
问题：{user_question}
思考过程：
"""

def self_consistency_reasoning(llm, prompt, n_samples=5):
    """自一致性推理：多次采样取多数结果"""
    answers = []
    
    for _ in range(n_samples):
        response = llm.generate(prompt, temperature=0.7)
        answer = extract_final_answer(response)
        answers.append(answer)
    
    # 投票选出最一致的答案
    from collections import Counter
    most_common = Counter(answers).most_common(1)[0]
    
    return {
        "answer": most_common[0],
        "confidence": most_common[1] / n_samples,
        "all_answers": answers
    }

class ContextManager:
    """上下文窗口管理器"""
    
    def __init__(self, max_tokens: int = 128000, 
                 reserve_tokens: int = 4096):
        self.max_tokens = max_tokens
        self.reserve_tokens = reserve_tokens  # 预留给输出的 token
        self.available_tokens = max_tokens - reserve_tokens
    
    def manage(self, system_prompt: str, messages: list, 
               tools: list = None) -> list:
        """管理上下文，确保不超出窗口"""
        # 计算固定部分的 token 数
        fixed_tokens = self._count_tokens(system_prompt)
        if tools:
            fixed_tokens += self._count_tokens(json.dumps(tools))
        
        available = self.available_tokens - fixed_tokens
        
        # 从最新消息开始，逆序添加
        managed_messages = []
        current_tokens = 0
        
        for msg in reversed(messages):
            msg_tokens = self._count_tokens(msg["content"])
            if current_tokens + msg_tokens > available:
                break
            managed_messages.insert(0, msg)
            current_tokens += msg_tokens
        
        # 如果丢弃了旧消息，添加摘要
        if len(managed_messages) < len(messages):
            dropped = messages[:len(messages) - len(managed_messages)]
            summary = self._summarize(dropped)
            managed_messages.insert(0, {
                "role": "system",
                "content": f"[历史对话摘要] {summary}"
            })
        
        return managed_messages

class TokenBudgetManager:
    """Token 预算管理器"""
    
    BUDGET_ALLOCATION = {
        "system_prompt": 0.10,
        "tool_definitions": 0.08,
        "history": 0.40,
        "retrieval": 0.15,
        "output_reserve": 0.15,
        "safety_margin": 0.12,
    }
    
    def __init__(self, max_tokens: int = 128000):
        self.max_tokens = max_tokens
        self.budgets = {
            k: int(v * max_tokens) 
            for k, v in self.BUDGET_ALLOCATION.items()
        }
    
    def allocate(self, system_tokens: int, tool_tokens: int) -> dict:
        """动态分配剩余预算"""
        used = system_tokens + tool_tokens
        remaining = self.max_tokens - used - self.budgets["output_reserve"]
        
        return {
            "history_budget": int(remaining * 0.65),
            "retrieval_budget": int(remaining * 0.25),
            "safety_margin": int(remaining * 0.10),
        }

class ModelRouter:
    """模型路由器：根据任务类型选择最优模型"""
    
    def __init__(self):
        self.models = {
            "reasoning": "deepseek-r1",        # 复杂推理
            "coding": "claude-3-5-sonnet",     # 代码生成
            "general": "gpt-4o",               # 通用任务
            "simple": "gpt-4o-mini",           # 简单任务
            "long_context": "gemini-1.5-pro"   # 长文档
        }
    
    def route(self, task: dict) -> str:
        """根据任务特征选择模型"""
        if task.get("requires_deep_reasoning"):
            return self.models["reasoning"]
        elif task.get("type") == "coding":
            return self.models["coding"]
        elif task.get("context_length", 0) > 100000:
            return self.models["long_context"]
        elif task.get("complexity") == "low":
            return self.models["simple"]
        else:
            return self.models["general"]

class ReActAgent:
    """ReAct 框架 Agent"""
    
    REACT_PROMPT = """你是一个能够使用工具解决问题的 AI Agent。

请按照以下格式进行推理和行动：

Thought: [分析当前情况，思考下一步该做什么]
Action: [选择要使用的工具]
Action Input: [工具的输入参数]
Observation: [工具返回的结果]
... (重复 Thought/Action/Observation 直到任务完成)
Thought: 我已经获得了足够的信息来回答问题
Final Answer: [最终答案]

可用工具：
{tools}

注意：
- 每次只执行一个 Action
- 仔细分析 Observation 再决定下一步
- 如果工具调用失败，分析原因并尝试其他方法
"""
    
    def __init__(self, llm, tools: dict, max_iterations: int = 10):
        self.llm = llm
        self.tools = tools
        self.max_iterations = max_iterations
    
    def run(self, task: str) -> str:
        """执行 ReAct 循环"""
        tool_descriptions = self._format_tools()
        system_prompt = self.REACT_PROMPT.format(tools=tool_descriptions)
        
        messages = [{"role": "user", "content": task}]
        
        for i in range(self.max_iterations):
            # LLM 生成 Thought + Action
            response = self.llm.generate(system_prompt, messages)
            messages.append({"role": "assistant", "content": response})
            
            # 检查是否有 Final Answer
            if "Final Answer:" in response:
                return self._extract_final_answer(response)
            
            # 解析 Action
            action, action_input = self._parse_action(response)
            
            if action is None:
                messages.append({
                    "role": "user", 
                    "content": "请按照格式输出 Action 和 Action Input"
                })
                continue
            
            # 执行 Action
            try:
                observation = self.tools[action].execute(action_input)
            except Exception as e:
                observation = f"工具执行错误: {str(e)}"
            
            # 将 Observation 反馈给 LLM
            messages.append({
                "role": "user",
                "content": f"Observation: {observation}"
            })
        
        return "达到最大迭代次数，任务未完成"

class PlanAndExecuteAgent:
    """Plan-and-Execute 框架 Agent"""
    
    def __init__(self, planner_llm, executor_llm, tools: dict):
        self.planner = planner_llm
        self.executor = executor_llm
        self.tools = tools
    
    def run(self, task: str) -> str:
        # 阶段 1：生成计划
        plan = self._create_plan(task)
        print(f"生成计划：{len(plan)} 个步骤")
        
        results = []
        
        # 阶段 2：逐步执行
        for i, step in enumerate(plan):
            print(f"执行步骤 {i+1}/{len(plan)}: {step['description']}")
            
            # 执行单个步骤
            result = self._execute_step(step, results)
            results.append({"step": step, "result": result})
            
            # 评估是否需要重规划
            if self._needs_replan(step, result, plan[i+1:]):
                remaining_plan = self._replan(task, results, plan[i+1:])
                plan = plan[:i+1] + remaining_plan
        
        # 阶段 3：汇总结果
        return self._summarize(task, results)
    
    def _create_plan(self, task: str) -> list:
        """使用 Planner LLM 生成执行计划"""
        prompt = f"""请为以下任务创建详细的执行计划。

任务：{task}

可用工具：{self._format_tools()}

请输出 JSON 格式的计划：
[
  {{"step": 1, "description": "...", "tool": "...", "expected_output": "..."}},
  ...
]"""
        
        response = self.planner.generate(prompt)
        return json.loads(response)
    
    def _execute_step(self, step: dict, previous_results: list) -> dict:
        """使用 Executor LLM 执行单个步骤"""
        context = self._build_context(previous_results)
        
        prompt = f"""请执行以下步骤：

步骤描述：{step['description']}
使用工具：{step['tool']}
前序结果：{context}

请生成工具调用参数并执行。"""
        
        # Executor 生成工具调用
        tool_call = self.executor.generate(prompt)
        
        # 执行工具
        tool = self.tools[step['tool']]
        return tool.execute(**json.loads(tool_call))

class ReflexionAgent:
    """Reflexion 反思框架 Agent"""
    
    def __init__(self, llm, tools, max_retries: int = 3):
        self.llm = llm
        self.tools = tools
        self.max_retries = max_retries
        self.reflections = []  # 反思记忆
    
    def run(self, task: str) -> str:
        for attempt in range(self.max_retries):
            # 执行任务（带上历史反思）
            result = self._execute_with_reflections(task)
            
            # 评估结果
            evaluation = self._evaluate(task, result)
            
            if evaluation["success"]:
                return result
            
            # 反思失败原因
            reflection = self._reflect(task, result, evaluation)
            self.reflections.append(reflection)
            print(f"第 {attempt+1} 次尝试失败，反思：{reflection}")
        
        return f"经过 {self.max_retries} 次尝试仍未成功。最后结果：{result}"
    
    def _reflect(self, task: str, result: str, evaluation: dict) -> str:
        """自我反思，分析失败原因"""
        prompt = f"""请分析以下任务执行失败的原因，并给出改进建议。

任务：{task}
执行结果：{result}
评估反馈：{evaluation['feedback']}

历史反思：
{self._format_reflections()}

请输出：
1. 失败原因分析
2. 具体改进建议
3. 下次执行时应该注意什么"""
        
        return self.llm.generate(prompt)
    
    def _execute_with_reflections(self, task: str) -> str:
        """带反思经验执行任务"""
        reflection_context = ""
        if self.reflections:
            reflection_context = f"""
从之前的尝试中学到的经验：
{self._format_reflections()}

请根据以上经验避免重复犯错。
"""
        
        prompt = f"""任务：{task}
{reflection_context}
请执行任务。"""
        
        return self._react_execute(prompt)

class TreeOfThoughtsAgent:
    """Tree of Thoughts 推理 Agent"""
    
    def __init__(self, llm, num_branches: int = 3, max_depth: int = 5):
        self.llm = llm
        self.num_branches = num_branches
        self.max_depth = max_depth
    
    def solve(self, problem: str) -> str:
        """使用 ToT 解决问题"""
        root = {"thought": "开始分析问题", "children": [], "score": 0}
        
        # BFS 搜索
        best_path = self._bfs_search(problem, root)
        
        return self._format_solution(best_path)
    
    def _bfs_search(self, problem: str, root: dict) -> list:
        """广度优先搜索最优推理路径"""
        queue = [(root, [root])]
        best_path = None
        best_score = -1
        
        for depth in range(self.max_depth):
            next_queue = []
            
            for node, path in queue:
                # 生成多个候选思路
                candidates = self._generate_thoughts(problem, path)
                
                for candidate in candidates:
                    # 评估每个思路的质量
                    score = self._evaluate_thought(problem, path, candidate)
                    candidate["score"] = score
                    
                    new_path = path + [candidate]
                    
                    # 检查是否到达终点
                    if self._is_solution(candidate):
                        if score > best_score:
                            best_score = score
                            best_path = new_path
                    else:
                        next_queue.append((candidate, new_path))
            
            # 剪枝：只保留得分最高的分支
            next_queue.sort(key=lambda x: x[0]["score"], reverse=True)
            queue = next_queue[:self.num_branches]
        
        return best_path
    
    def _generate_thoughts(self, problem: str, path: list) -> list:
        """生成多个候选思路"""
        context = "\n".join([n["thought"] for n in path])
        
        prompt = f"""问题：{problem}

当前推理路径：
{context}

请生成 {self.num_branches} 个不同的下一步思路。
每个思路要有不同的角度和方法。

格式：
思路1: [描述]
思路2: [描述]
思路3: [描述]"""
        
        response = self.llm.generate(prompt)
        return self._parse_thoughts(response)
    
    def _evaluate_thought(self, problem: str, path: list, 
                          thought: dict) -> float:
        """评估思路的质量（0-1 分）"""
        prompt = f"""评估以下推理步骤的质量。

问题：{problem}
推理路径：{[n['thought'] for n in path]}
当前步骤：{thought['thought']}

请从以下维度评分（0-1）：
1. 逻辑正确性
2. 与问题的相关性
3. 推进问题解决的程度
4. 是否引入新的有效信息

总分（0-1）："""
        
        response = self.llm.generate(prompt)
        return float(response.strip())

def ucb1_score(node, parent, exploration_weight=1.414):
    """UCB1 选择公式"""
    if node.visits == 0:
        return float('inf')
    
    exploitation = node.value / node.visits
    exploration = exploration_weight * math.sqrt(
        math.log(parent.visits) / node.visits
    )
    return exploitation + exploration

class ReasoningOptimizer:
    """推理过程优化器"""
    
    def __init__(self, max_steps=10):
        self.max_steps = max_steps
        self.action_history = []
    
    def check_loop(self, current_action: dict) -> bool:
        """检测推理循环"""
        # 检查最近 3 步是否有重复
        recent = self.action_history[-3:]
        for past_action in recent:
            if self._is_similar(current_action, past_action):
                return True  # 检测到循环
        return False
    
    def check_relevance(self, current_thought: str, 
                        original_goal: str) -> float:
        """检查当前推理与目标的相关性"""
        # 使用 Embedding 计算相似度
        goal_emb = self.embed(original_goal)
        thought_emb = self.embed(current_thought)
        return cosine_similarity(goal_emb, thought_emb)
    
    def should_stop(self, step: int, relevance: float) -> bool:
        """判断是否应该停止推理"""
        if step >= self.max_steps:
            return True
        if relevance < 0.3:  # 严重偏离目标
            return True
        return False

import openai

# 定义工具
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称，如：北京、上海"
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "温度单位"
                    }
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "搜索互联网获取最新信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "搜索关键词"
                    },
                    "max_results": {
                        "type": "integer",
                        "description": "最大结果数",
                        "default": 5
                    }
                },
                "required": ["query"]
            }
        }
    }
]

# 调用 LLM
response = openai.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "user", "content": "北京今天天气怎么样？"}
    ],
    tools=tools,
    tool_choice="auto"  # 让模型自动决定是否调用工具
)

# 解析工具调用
if response.choices[0].message.tool_calls:
    tool_call = response.choices[0].message.tool_calls[0]
    print(f"调用工具: {tool_call.function.name}")
    print(f"参数: {tool_call.function.arguments}")

from typing import Callable, Any
from dataclasses import dataclass, field
import inspect

@dataclass
class ToolDefinition:
    """工具定义"""
    name: str
    description: str
    function: Callable
    parameters: dict
    examples: list = field(default_factory=list)
    category: str = "general"
    requires_confirmation: bool = False  # 是否需要用户确认

class ToolRegistry:
    """工具注册中心"""
    
    def __init__(self):
        self._tools: dict[str, ToolDefinition] = {}
    
    def register(self, name: str = None, description: str = None,
                 category: str = "general", 
                 requires_confirmation: bool = False):
        """装饰器方式注册工具"""
        def decorator(func: Callable):
            tool_name = name or func.__name__
            tool_desc = description or func.__doc__ or ""
            
            # 自动从函数签名提取参数定义
            params = self._extract_parameters(func)
            
            self._tools[tool_name] = ToolDefinition(
                name=tool_name,
                description=tool_desc,
                function=func,
                parameters=params,
                category=category,
                requires_confirmation=requires_confirmation
            )
            return func
        return decorator
    
    def _extract_parameters(self, func: Callable) -> dict:
        """从函数签名自动提取参数定义"""
        sig = inspect.signature(func)
        hints = func.__annotations__
        
        properties = {}
        required = []
        
        for param_name, param in sig.parameters.items():
            if param_name == "self":
                continue
            
            param_type = hints.get(param_name, str).__name__
            type_map = {"str": "string", "int": "integer", 
                       "float": "number", "bool": "boolean"}
            
            properties[param_name] = {
                "type": type_map.get(param_type, "string"),
                "description": f"参数 {param_name}"
            }
            
            if param.default == inspect.Parameter.empty:
                required.append(param_name)
        
        return {
            "type": "object",
            "properties": properties,
            "required": required
        }
    
    def to_openai_tools(self) -> list:
        """转换为 OpenAI Function Calling 格式"""
        return [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters
                }
            }
            for tool in self._tools.values()
        ]
    
    def execute(self, name: str, **kwargs) -> Any:
        """执行工具"""
        tool = self._tools.get(name)
        if not tool:
            raise ValueError(f"工具 {name} 未注册")
        return tool.function(**kwargs)


# 使用示例
registry = ToolRegistry()

@registry.register(
    description="执行 SQL 查询并返回结果",
    category="database",
    requires_confirmation=True
)
def query_database(sql: str, database: str = "default") -> str:
    """执行 SQL 查询"""
    # 实际执行逻辑
    pass

@registry.register(
    description="发送 HTTP 请求",
    category="network"
)
def http_request(url: str, method: str = "GET", body: str = None) -> str:
    """发送 HTTP 请求"""
    pass

class ToolExecutionEngine:
    """工具调用引擎"""
    
    def __init__(self, registry: ToolRegistry, max_retries: int = 2):
        self.registry = registry
        self.max_retries = max_retries
        self.execution_log = []
    
    async def execute_tool_calls(self, tool_calls: list) -> list:
        """批量执行工具调用（支持并行）"""
        import asyncio
        
        tasks = [
            self._execute_single(tc) for tc in tool_calls
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        return [
            {
                "tool_call_id": tc.get("id"),
                "name": tc["name"],
                "result": str(r) if not isinstance(r, Exception) else f"Error: {r}",
                "status": "success" if not isinstance(r, Exception) else "error"
            }
            for tc, r in zip(tool_calls, results)
        ]
    
    async def _execute_single(self, tool_call: dict) -> Any:
        """执行单个工具调用（带重试）"""
        name = tool_call["name"]
        args = tool_call.get("arguments", {})
        
        for attempt in range(self.max_retries + 1):
            try:
                result = self.registry.execute(name, **args)
                
                # 记录执行日志
                self.execution_log.append({
                    "tool": name,
                    "args": args,
                    "result": str(result)[:500],
                    "attempt": attempt + 1,
                    "status": "success"
                })
                
                return result
                
            except Exception as e:
                if attempt < self.max_retries:
                    await asyncio.sleep(1 * (attempt + 1))  # 退避重试
                    continue
                
                self.execution_log.append({
                    "tool": name,
                    "args": args,
                    "error": str(e),
                    "attempt": attempt + 1,
                    "status": "failed"
                })
                raise

class ToolOrchestrator:
    """工具编排器"""
    
    def __init__(self, registry: ToolRegistry):
        self.registry = registry
    
    async def sequential(self, steps: list) -> list:
        """顺序执行"""
        results = []
        context = {}
        
        for step in steps:
            # 将前序结果注入参数
            args = self._inject_context(step["args"], context)
            result = self.registry.execute(step["tool"], **args)
            results.append(result)
            context[step.get("output_key", f"step_{len(results)}")] = result
        
        return results
    
    async def parallel(self, steps: list) -> list:
        """并行执行"""
        import asyncio
        
        tasks = [
            asyncio.to_thread(self.registry.execute, step["tool"], **step["args"])
            for step in steps
        ]
        return await asyncio.gather(*tasks)
    
    async def pipeline(self, steps: list, initial_input: Any) -> Any:
        """管道执行：前一步输出作为后一步输入"""
        current = initial_input
        
        for step in steps:
            args = {step.get("input_key", "input"): current, **step.get("extra_args", {})}
            current = self.registry.execute(step["tool"], **args)
        
        return current

class ToolResultProcessor:
    """工具结果后处理器"""
    
    MAX_LENGTH = 2000  # 最大返回长度
    
    def process(self, tool_name: str, raw_result: Any) -> str:
        """处理工具返回结果"""
        result_str = str(raw_result)
        
        # 1. 长度控制
        if len(result_str) > self.MAX_LENGTH:
            result_str = self._truncate_smart(result_str)
        
        # 2. 格式优化
        if isinstance(raw_result, dict):
            result_str = self._format_dict(raw_result)
        elif isinstance(raw_result, list):
            result_str = self._format_list(raw_result)
        
        # 3. 错误处理
        if "error" in result_str.lower():
            result_str = self._humanize_error(result_str)
        
        return result_str
    
    def _truncate_smart(self, text: str) -> str:
        """智能截断：保留开头和结尾"""
        head = text[:self.MAX_LENGTH // 2]
        tail = text[-self.MAX_LENGTH // 4:]
        return f"{head}\n\n...(内容已截断，共 {len(text)} 字符)...\n\n{tail}"

class ShortTermMemory:
    """短期记忆管理"""
    
    def __init__(self, max_tokens: int = 100000):
        self.messages = []
        self.max_tokens = max_tokens
        self.token_counter = TokenCounter()
    
    def add(self, role: str, content: str, importance: float = 0.5):
        """添加消息到短期记忆"""
        self.messages.append({
            "role": role,
            "content": content,
            "importance": importance,
            "timestamp": datetime.now(),
            "tokens": self.token_counter.count(content)
        })
        
        # 检查是否超出限制
        self._manage_capacity()
    
    def _manage_capacity(self):
        """管理记忆容量"""
        total_tokens = sum(m["tokens"] for m in self.messages)
        
        while total_tokens > self.max_tokens and len(self.messages) > 2:
            # 策略 1：移除最旧且最不重要的消息
            candidates = self.messages[1:-1]  # 保留第一条和最后一条
            least_important = min(candidates, 
                                  key=lambda m: m["importance"])
            
            # 策略 2：将被移除的消息摘要化
            summary = self._summarize_message(least_important)
            
            idx = self.messages.index(least_important)
            self.messages[idx] = {
                "role": "system",
                "content": f"[摘要] {summary}",
                "importance": 0.3,
                "timestamp": least_important["timestamp"],
                "tokens": self.token_counter.count(summary)
            }
            
            total_tokens = sum(m["tokens"] for m in self.messages)
    
    def get_context(self) -> list:
        """获取当前上下文"""
        return [{"role": m["role"], "content": m["content"]} 
                for m in self.messages]