from typing import Callable, Dict, Any
import json

class Tool:
    def __init__(self, name: str, description: str, func: Callable):
        self.name = name
        self.description = description
        self.func = func
    
    def run(self, **kwargs) -> str:
        try:
            return str(self.func(**kwargs))
        except Exception as e:
            return f"Error: {str(e)}"

class SimpleAgent:
    def __init__(self, model: str = "gpt-4"):
        self.model = model
        self.tools: Dict[str, Tool] = {}
        self.history = []
        self.max_iterations = 10
    
    def register_tool(self, tool: Tool):
        self.tools[tool.name] = tool
    
    def _build_system_prompt(self) -> str:
        tool_descriptions = "\n".join([
            f"- {name}: {tool.description}"
            for name, tool in self.tools.items()
        ])
        
        return f"""你是一个有帮助的 AI 助手。你可以使用以下工具来完成任务：

{tool_descriptions}

当你需要使用工具时，请严格按照以下 JSON 格式输出：
{{
    "thought": "你思考的内容",
    "action": "工具名称",
    "action_input": {{工具参数字典}}
}}

如果你认为任务已经完成，不需要再调用工具，请按照以下格式输出：
{{
    "thought": "总结一下你做了什么",
    "final_answer": "给用户的最终回复"
}}

重要：每次只能调用一个工具，不要同时调用多个。
"""
    
    def run(self, user_input: str) -> str:
        self.history.append({"role": "user", "content": user_input})
        
        for i in range(self.max_iterations):
            # 调用大模型获取决策
            response = self._call_llm()
            
            try:
                result = json.loads(response)
            except json.JSONDecodeError:
                # 解析失败，重试
                self.history.append({
                    "role": "assistant", 
                    "content": "输出格式错误，请重新输出正确的 JSON"
                })
                continue
            
            # 检查是否给出了最终答案
            if "final_answer" in result:
                return result["final_answer"]
            
            # 调用工具
            tool_name = result["action"]
            if tool_name not in self.tools:
                observation = f"错误：工具 {tool_name} 不存在"
            else:
                tool = self.tools[tool_name]
                observation = tool.run(**result["action_input"])
            
            # 将思考、行动和观察加入历史
            self.history.append({
                "role": "assistant",
                "content": json.dumps({
                    "thought": result["thought"],
                    "action": tool_name,
                    "action_input": result["action_input"]
                }, ensure_ascii=False)
            })
            self.history.append({
                "role": "user",
                "content": f"观察结果：{observation}"
            })
        
        return "执行次数超出限制，任务未完成"

import subprocess
import textwrap

def execute_python(code: str) -> str:
    """执行 Python 代码并返回输出"""
    # 安全限制：禁止危险操作（生产环境需要更严格的沙箱）
    dangerous = ["os.system", "subprocess", "eval", "exec", "__import__"]
    for d in dangerous:
        if d in code:
            return f"安全限制：不允许使用 {d}"
    
    try:
        result = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=10
        )
        output = result.stdout
        if result.stderr:
            output += "\n错误：" + result.stderr
        return output if output else "代码执行成功，无输出"
    except subprocess.TimeoutExpired:
        return "执行超时（超过10秒）"
    except Exception as e:
        return f"执行错误：{str(e)}"

python_tool = Tool(
    name="execute_python",
    description="执行 Python 代码，返回标准输出。适用于计算、数据处理、文件操作等任务。",
    func=execute_python
)

import requests
from bs4 import BeautifulSoup

def web_search(query: str) -> str:
    """搜索网络信息"""
    try:
        # 使用 DuckDuckGo 搜索（不需要 API Key）
        url = f"https://html.duckduckgo.com/html/?q={query}"
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(url, headers=headers, timeout=10)
        soup = BeautifulSoup(response.text, "html.parser")
        
        results = []
        for result in soup.select(".result")[:5]:
            title = result.select_one(".result__a").get_text()
            snippet = result.select_one(".result__snippet").get_text()
            results.append(f"标题：{title}\n摘要：{snippet}\n")
        
        return "\n".join(results) if results else "没有找到搜索结果"
    except Exception as e:
        return f"搜索失败：{str(e)}"

search_tool = Tool(
    name="web_search",
    description="搜索网络信息，获取最新的新闻、技术文档、价格等实时数据。",
    func=web_search
)

agent = SimpleAgent()
agent.register_tool(python_tool)
agent.register_tool(search_tool)

# 测试：让 Agent 计算斐波那契数列的第30项
result = agent.run("帮我计算斐波那契数列的第30项是多少")
print(result)

from collections import deque
from typing import List, Dict

class HierarchicalMemory:
    def __init__(self, max_short_term: int = 5, max_mid_term: int = 10):
        self.short_term = deque(maxlen=max_short_term)
        self.mid_term = []
        self.long_term = []  # 实际项目中换成向量数据库
    
    def add(self, message: Dict[str, str]):
        self.short_term.append(message)
        # 短期记忆满了，将最早的一条摘要后移入中期记忆
        if len(self.short_term) == self.short_term.maxlen:
            oldest = self.short_term[0]
            summarized = self._summarize(oldest)
            self.mid_term.append(summarized)
            if len(self.mid_term) > max_mid_term:
                # 中期记忆满了，压缩成长期记忆
                self._compress_to_long_term()
    
    def _summarize(self, message: Dict[str, str]) -> Dict[str, str]:
        """对消息进行摘要（实际项目中调用LLM摘要）"""
        content = message["content"]
        if len(content) > 200:
            content = content[:200] + "..."
        return {"role": message["role"], "content": content, "summarized": True}
    
    def _compress_to_long_term(self):
        """将中期记忆压缩为长期记忆"""
        summary = f"历史执行摘要：共{len(self.mid_term)}条记录"
        self.long_term.append(summary)
        self.mid_term = self.mid_term[-3:]  # 保留最近3条
    
    def get_context(self) -> List[Dict[str, str]]:
        """获取当前上下文中的所有消息"""
        context = []
        if self.long_term:
            context.append({
                "role": "system",
                "content": "【长期记忆】" + "\n".join(self.long_term)
            })
        context.extend(self.mid_term)
        context.extend(self.short_term)
        return context

import time
from functools import wraps

def retry(max_attempts: int = 3, delay: float = 1.0):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_error = None
            for attempt in range(max_attempts):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_error = e
                    if attempt < max_attempts - 1:
                        time.sleep(delay * (2 ** attempt))  # 指数退避
            raise last_error
        return wrapper
    return decorator

class RobustAgent(SimpleAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.max_tool_retries = 3
    
    def _execute_tool(self, tool_name: str, action_input: Dict) -> str:
        for attempt in range(self.max_tool_retries):
            try:
                tool = self.tools[tool_name]
                return tool.run(**action_input)
            except KeyError:
                return f"错误：工具 {tool_name} 不存在，可用工具：{list(self.tools.keys())}"
            except TypeError as e:
                return f"参数错误：{str(e)}。请检查参数是否匹配工具要求。"
            except Exception as e:
                if attempt < self.max_tool_retries - 1:
                    time.sleep(1)
                    continue
                return f"工具执行失败（已重试{self.max_tool_retries}次）：{str(e)}"

class ReflectiveAgent(RobustAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.reflection_enabled = True
    
    def _reflect(self) -> str:
        """让 Agent 反思当前的执行过程"""
        reflection_prompt = """
请反思你刚才的执行过程，回答以下问题：
1. 刚才的工具调用是否成功？
2. 如果失败了，失败的根本原因是什么？
3. 接下来应该如何调整策略？
4. 是否应该继续尝试，还是换一种方法？

请用简洁的语言回答。
"""
        return self._call_llm_with_prompt(reflection_prompt)
    
    def run(self, user_input: str) -> str:
        self.history.append({"role": "user", "content": user_input})
        
        for i in range(self.max_iterations):
            response = self._call_llm()
            result = json.loads(response)
            
            if "final_answer" in result:
                return result["final_answer"]
            
            # 执行工具
            observation = self._execute_tool(
                result["action"], 
                result["action_input"]
            )
            
            # 记录行动和观察
            self.history.append({
                "role": "assistant",
                "content": json.dumps(result, ensure_ascii=False)
            })
            self.history.append({
                "role": "user",
                "content": f"观察结果：{observation}"
            })
            
            # 如果出现错误，进行反思
            if self.reflection_enabled and "错误" in observation or "Error" in observation:
                reflection = self._reflect()
                self.history.append({
                    "role": "system",
                    "content": f"反思：{reflection}"
                })
        
        return "执行次数超出限制"

class CostAwareAgent(ReflectiveAgent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.total_tokens = 0
        self.max_tokens = 50000  # 单次任务最大 token 消耗
        self.token_price = 0.002 / 1000  # GPT-4的价格（示例）
    
    def _call_llm(self) -> str:
        # 检查预算
        if self.total_tokens >= self.max_tokens:
            raise Exception(f"Token 消耗已达上限 {self.max_tokens}，任务暂停")
        
        response = super()._call_llm()
        
        # 估算 token 消耗（实际应该用 tiktoken 精确计算）
        tokens_used = len(json.dumps(self.history)) // 4
        self.total_tokens += tokens_used
        
        estimated_cost = self.total_tokens * self.token_price
        if estimated_cost > 0.5:  # 超过0.5美元时提醒
            print(f"当前预估成本：${estimated_cost:.4f}")
        
        return response

class MultiAgentSystem:
    def __init__(self):
        self.agents = {}
        self.task_queue = []
        self.results = {}
    
    def register_agent(self, name: str, agent: SimpleAgent, role: str):
        self.agents[name] = {
            "agent": agent,
            "role": role,
            "status": "idle"
        }
    
    def assign_task(self, agent_name: str, task: str):
        if agent_name not in self.agents:
            raise ValueError(f"Agent {agent_name} 不存在")
        
        self.agents[agent_name]["status"] = "working"
        result = self.agents[agent_name]["agent"].run(task)
        self.agents[agent_name]["status"] = "idle"
        self.results[agent_name] = result
        return result
    
    def run_software_project(self, requirement: str):
        print("=" * 50)
        print("开始执行软件开发项目")
        print(f"原始需求：{requirement}")
        print("=" * 50)
        
        # 第一步：产品经理分析需求
        print("\n【产品经理 Agent 开始工作】")
        prd = self.assign_task(
            "product_manager",
            f"""请分析以下需求，输出产品需求文档：
需求：{requirement}

请包含以下内容：
1. 功能需求列表
2. 技术实现要点
3. 验收标准
"""
        )
        print("需求分析完成")
        print(prd[:500] + "..." if len(prd) > 500 else prd)
        
        # 第二步：程序员实现代码
        print("\n【程序员 Agent 开始工作】")
        code = self.assign_task(
            "programmer",
            f"""请根据以下产品需求编写 Python 代码：
{prd}

要求：
1. 代码结构清晰，有注释
2. 包含必要的错误处理
3. 提供使用示例
"""
        )
        print("代码编写完成")
        print(code[:500] + "..." if len(code) > 500 else code)
        
        # 第三步：测试工程师测试代码
        print("\n【测试工程师 Agent 开始工作】")
        test_report = self.assign_task(
            "tester",
            f"""请测试以下代码，输出测试报告：
代码：
{code}

请检查：
1. 代码是否有语法错误
2. 逻辑是否正确
3. 边界情况是否处理
4. 给出修复建议
"""
        )
        print("测试完成")
        print(test_report[:500] + "..." if len(test_report) > 500 else test_report)
        
        return {
            "prd": prd,
            "code": code,
            "test_report": test_report
        }

# 初始化各个 Agent
system = MultiAgentSystem()

system.register_agent(
    "product_manager",
    SimpleAgent(model="gpt-4"),
    role="产品经理"
)

system.register_agent(
    "programmer",
    SimpleAgent(model="gpt-4"),
    role="程序员"
)

system.register_agent(
    "tester",
    SimpleAgent(model="gpt-4"),
    role="测试工程师"
)

# 运行项目
result = system.run_software_project(
    "写一个命令行工具，能够统计指定目录下各种文件类型的数量和总大小"
)

import os
import ast
from pathlib import Path

def list_files(directory: str, pattern: str = "*.py") -> str:
    """列出目录下的所有文件"""
    try:
        files = []
        for root, dirs, filenames in os.walk(directory):
            for filename in filenames:
                if filename.endswith(pattern.replace("*", "")):
                    filepath = os.path.join(root, filename)
                    size = os.path.getsize(filepath)
                    files.append(f"{filepath} ({size} bytes)")
        return "\n".join(files) if files else "没有找到文件"
    except Exception as e:
        return f"错误：{str(e)}"

def read_file(filepath: str, max_lines: int = 100) -> str:
    """读取文件内容"""
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            lines = f.readlines()
            content = "".join(lines[:max_lines])
            if len(lines) > max_lines:
                content += f"\n...（文件共{len(lines)}行，已显示前{max_lines}行）"
            return content
    except Exception as e:
        return f"错误：{str(e)}"

def analyze_python_structure(filepath: str) -> str:
    """分析 Python 代码的结构（类、函数、导入）"""
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            source = f.read()
        
        tree = ast.parse(source)
        
        classes = []
        functions = []
        imports = []
        
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.append(alias.name)
            elif isinstance(node, ast.ImportFrom):
                imports.append(f"{node.module}")
            elif isinstance(node, ast.ClassDef):
                methods = [n.name for n in node.body if isinstance(n, ast.FunctionDef)]
                classes.append(f"{node.name}（方法：{', '.join(methods)}）")
            elif isinstance(node, ast.FunctionDef):
                args = [a.arg for a in node.args.args]
                functions.append(f"{node.name}（参数：{', '.join(args)}）")
        
        result = f"文件：{filepath}\n\n"
        result += f"导入的模块：\n- {chr(10)}- ".join(imports[:20]) if imports else "无"
        result += "\n\n类定义：\n"
        result += "\n".join([f"- {c}" for c in classes]) if classes else "- 无"
        result += "\n\n函数定义：\n"
        result += "\n".join([f"- {f}" for f in functions]) if functions else "- 无"
        
        return result
    except Exception as e:
        return f"分析失败：{str(e)}"

def write_documentation(content: str, output_path: str) -> str:
    """将生成的文档写入文件"""
    try:
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(content)
        return f"文档已写入：{output_path}"
    except Exception as e:
        return f"写入失败：{str(e)}"

# 注册工具
list_tool = Tool("list_files", "列出目录下的文件，支持通配符", list_files)
read_tool = Tool("read_file", "读取文件内容，可指定最大行数", read_file)
analyze_tool = Tool("analyze_python", "分析 Python 代码的结构", analyze_python_structure)
write_tool = Tool("write_doc", "将生成的文档写入文件", write_documentation)

class DocumentationAgent:
    def __init__(self):
        self.agent = SimpleAgent(model="gpt-4")
        self.agent.register_tool(list_tool)
        self.agent.register_tool(read_tool)
        self.agent.register_tool(analyze_tool)
        self.agent.register_tool(write_tool)
        
        self.project_info = {}
    
    def generate_docs(self, project_path: str, output_path: str):
        print(f"开始分析项目：{project_path}")
        
        # 第一步：列出所有文件
        files_result = self.agent.run(
            f"请列出 {project_path} 目录下所有的 Python 文件，不要递归太深"
        )
        print("文件列表获取完成")
        
        # 第二步：分析每个核心文件的结构
        structure_result = self.agent.run("""
请分析主要的 Python 文件，提取以下信息：
1. 每个文件的主要功能
2. 定义的类和方法
3. 主要的函数接口
4. 模块之间的依赖关系

只分析核心文件，忽略测试文件和示例文件。
""")
        print("代码结构分析完成")
        
        # 第三步：生成架构概述
        architecture_result = self.agent.run("""
基于以上分析，请生成项目的架构概述文档，包括：
1. 项目的整体架构设计
2. 核心模块的职责说明
3. 关键的设计决策
4. 模块之间的交互关系

用清晰的 Markdown 格式输出。
""")
        print("架构概述生成完成")
        
        # 第四步：生成 API 文档
        api_result = self.agent.run("""
现在请生成详细的 API 文档，包括每个主要类和函数的：
1. 功能说明
2. 参数说明
3. 返回值说明
4. 使用示例（如果能推断出来）

用 Markdown 格式组织，结构清晰。
""")
        print("API 文档生成完成")
        
        # 第五步：生成使用指南
        usage_result = self.agent.run("""
最后生成项目的使用指南，包括：
1. 安装步骤
2. 快速开始示例
3. 常见使用场景
4. 注意事项和最佳实践

确保文档对新用户友好。
""")
        print("使用指南生成完成")
        
        # 合并所有文档并写入
        final_doc = self.agent.run(f"""
请将以上生成的所有文档内容整合成一份完整的技术文档，结构如下：

# 项目技术文档

## 1. 项目概述
## 2. 架构设计
## 3. API 参考
## 4. 使用指南
## 5. 注意事项

然后调用 write_doc 工具将文档写入 {output_path}。
""")
        
        print(f"文档生成完成：{output_path}")
        return final_doc

doc_agent = DocumentationAgent()
result = doc_agent.generate_docs(
    project_path="/path/to/your/project",
    output_path="/path/to/output/docs.md"
)