# agent/tools.py
import requests
import json
from typing import Any

# 工具函数：网页搜索
def web_search(query: str, max_results: int = 3) -> list[dict]:
    """
    使用搜索引擎查询信息。
    
    参数:
        query: 搜索关键词
        max_results: 最大返回结果数
    
    返回:
        搜索结果列表，每项包含 title, url, snippet
    """
    # 这里使用一个简单的示例 API，实际使用时替换为你的搜索 API
    url = f"https://api.duckduckgo.com/?q={query}&format=json"
    try:
        response = requests.get(url, timeout=10)
        data = response.json()
        results = []
        for item in data.get("results", [])[:max_results]:
            results.append({
                "title": item.get("Title", ""),
                "url": item.get("FirstURL", ""),
                "snippet": item.get("Abstract", "")
            })
        return results
    except Exception as e:
        return [{"error": f"搜索失败: {str(e)}"}]

# 工具函数：计算器
def calculator(expression: str) -> float | str:
    """
    计算数学表达式。
    
    参数:
        expression: 数学表达式字符串，如 "2 + 3 * 4"
    
    返回:
        计算结果
    """
    try:
        # 安全计算：仅允许基本的数学运算
        allowed = set("0123456789+-*/()., ")
        if not all(c in allowed for c in expression):
            return "错误：表达式包含不允许的字符"
        result = eval(expression, {"__builtins__": {}}, {})
        return result
    except Exception as e:
        return f"计算错误: {str(e)}"

# 工具注册表
TOOL_REGISTRY = {
    "web_search": {
        "function": web_search,
        "description": "搜索互联网获取最新信息",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "搜索关键词"
                },
                "max_results": {
                    "type": "integer",
                    "description": "最大返回结果数（默认 3）",
                    "default": 3
                }
            },
            "required": ["query"]
        }
    },
    "calculator": {
        "function": calculator,
        "description": "计算数学表达式",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "数学表达式"
                }
            },
            "required": ["expression"]
        }
    }
}

# agent/core.py
import json
import os
from openai import OpenAI
from dotenv import load_dotenv
from .tools import TOOL_REGISTRY

load_dotenv()

class SimpleAgent:
    """一个简单的 AI Agent 实现"""
    
    def __init__(self, model: str = "gpt-4o"):
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.model = model
        self.messages = []  # 对话历史
        self.tools = self._build_tools_schema()
    
    def _build_tools_schema(self) -> list[dict]:
        """将工具注册表转换为 OpenAI 工具调用格式"""
        tools = []
        for name, meta in TOOL_REGISTRY.items():
            tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": meta["description"],
                    "parameters": meta["parameters"]
                }
            })
        return tools
    
    def _execute_tool(self, name: str, arguments: dict) -> str:
        """执行工具调用"""
        if name not in TOOL_REGISTRY:
            return f"错误：未知工具 '{name}'"
        try:
            result = TOOL_REGISTRY[name]["function"](**arguments)
            return json.dumps(result, ensure_ascii=False)
        except Exception as e:
            return f"工具执行错误: {str(e)}"
    
    def run(self, user_input: str, max_iterations: int = 5) -> str:
        """
        运行 Agent，处理用户输入。
        
        Perceive → Think → Act → Observe 循环
        """
        self.messages.append({"role": "user", "content": user_input})
        
        for iteration in range(max_iterations):
            print(f"\n[迭代 {iteration + 1}/{max_iterations}]")
            
            # 1. Perceive + Think: 发送给 LLM
            response = self.client.chat.completions.create(
                model=self.model,
                messages=self.messages,
                tools=self.tools,
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            
            # 检查是否需要工具调用
            if message.tool_calls:
                # 2. Act: 执行工具
                self.messages.append(message)
                for tool_call in message.tool_calls:
                    func_name = tool_call.function.name
                    func_args = json.loads(tool_call.function.arguments)
                    print(f"  → 调用工具: {func_name}({func_args})")
                    
                    # 3. Observe: 观察结果
                    result = self._execute_tool(func_name, func_args)
                    print(f"  ← 结果: {result[:100]}...")
                    
                    self.messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": result
                    })
            else:
                # 不需要工具调用，直接返回模型回答
                final_answer = message.content
                self.messages.append({"role": "assistant", "content": final_answer})
                return final_answer
        
        # 达到最大迭代次数，返回当前结果
        return "Agent 已达到最大迭代次数，请尝试简化问题。"
    
    def reset(self):
        """重置对话历史"""
        self.messages = []

# main.py
from agent.core import SimpleAgent

def main():
    agent = SimpleAgent()
    
    print("=" * 60)
    print("  AI Agent 演示 - 输入 'exit' 退出")
    print("=" * 60)
    
    while True:
        user_input = input("\n👤 你: ").strip()
        if user_input.lower() in ("exit", "quit", "q"):
            print("再见！")
            break
        
        print("\n🤖 Agent 思考中...")
        result = agent.run(user_input)
        print(f"\n🤖 Agent: {result}")

if __name__ == "__main__":
    main()

# agent/tools.py（扩展）

# 工具：文件读取器
def read_file_content(filepath: str, max_length: int = 2000) -> str:
    """
    读取本地文件内容。
    
    参数:
        filepath: 文件路径
        max_length: 最大读取字符数
    
    返回:
        文件内容
    """
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read(max_length)
        return content
    except Exception as e:
        return f"读取文件失败: {str(e)}"

# 工具：获取当前时间
def get_current_time() -> dict:
    """
    获取当前日期和时间。
    
    返回:
        包含日期、时间、星期几的字典
    """
    from datetime import datetime
    now = datetime.now()
    weekdays = ["周一", "周二", "周三", "周四", "周五", "周六", "周日"]
    return {
        "date": now.strftime("%Y-%m-%d"),
        "time": now.strftime("%H:%M:%S"),
        "weekday": weekdays[now.weekday()]
    }

# 工具：URL 内容抓取
def fetch_webpage(url: str) -> str:
    """
    获取网页的文本内容。
    
    参数:
        url: 网页 URL
    
    返回:
        网页文本内容
    """
    try:
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
        response = requests.get(url, headers=headers, timeout=15)
        response.raise_for_status()
        # 简单提取文本（实际项目中可用 BeautifulSoup）
        text = response.text[:3000]
        return text
    except Exception as e:
        return f"抓取失败: {str(e)}"

# 在 TOOL_REGISTRY 中添加新条目
TOOL_REGISTRY["read_file"] = {
    "function": read_file_content,
    "description": "读取本地文件的内容",
    "parameters": {
        "type": "object",
        "properties": {
            "filepath": {
                "type": "string",
                "description": "文件路径"
            },
            "max_length": {
                "type": "integer",
                "description": "最大读取字符数",
                "default": 2000
            }
        },
        "required": ["filepath"]
    }
}

TOOL_REGISTRY["get_current_time"] = {
    "function": get_current_time,
    "description": "获取当前日期和时间",
    "parameters": {
        "type": "object",
        "properties": {}
    }
}

TOOL_REGISTRY["fetch_webpage"] = {
    "function": fetch_webpage,
    "description": "获取指定 URL 的网页文本内容",
    "parameters": {
        "type": "object",
        "properties": {
            "url": {
                "type": "string",
                "description": "网页 URL"
            }
        },
        "required": ["url"]
    }
}

# agent/memory.py
import json
import os
from datetime import datetime
from typing import Optional

class ConversationMemory:
    """管理对话历史（短期记忆）"""
    
    def __init__(self, max_context_length: int = 4096):
        self.messages = []
        self.max_context_length = max_context_length  # 最大 token 数（近似）
        self.current_length = 0
    
    def add_message(self, role: str, content: str):
        """添加消息到对话历史"""
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }
        self.messages.append(message)
        self.current_length += len(content)
        self._prune_context()
    
    def get_context(self) -> list[dict]:
        """获取当前上下文（用于 LLM 调用）"""
        return [
            {"role": msg["role"], "content": msg["content"]}
            for msg in self.messages
        ]
    
    def _prune_context(self):
        """当上下文超过限制时，裁剪最早的消息"""
        while self.current_length > self.max_context_length and len(self.messages) > 1:
            removed = self.messages.pop(0)
            self.current_length -= len(removed["content"])
    
    def clear(self):
        """清空对话历史"""
        self.messages = []
        self.current_length = 0

# agent/memory.py（续）

class PersistentMemory:
    """持久化记忆（长期记忆）—— 保存到本地文件"""
    
    def __init__(self, storage_path: str = "./agent_memory.json"):
        self.storage_path = storage_path
        self.memories = self._load()
    
    def _load(self) -> dict:
        """从文件加载记忆"""
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, "r", encoding="utf-8") as f:
                    return json.load(f)
            except (json.JSONDecodeError, IOError):
                return {"sessions": [], "knowledge": {}}
        return {"sessions": [], "knowledge": {}}
    
    def _save(self):
        """保存记忆到文件"""
        with open(self.storage_path, "w", encoding="utf-8") as f:
            json.dump(self.memories, f, ensure_ascii=False, indent=2)
    
    def save_session(self, session_id: str, messages: list[dict]):
        """保存一次对话会话"""
        session = {
            "id": session_id,
            "timestamp": datetime.now().isoformat(),
            "message_count": len(messages),
            "messages": messages[-10:]  # 只保存最后 10 条
        }
        self.memories["sessions"].append(session)
        # 最多保留 50 个会话
        if len(self.memories["sessions"]) > 50:
            self.memories["sessions"] = self.memories["sessions"][-50:]
        self._save()
    
    def store_knowledge(self, key: str, value: str):
        """存储知识片段"""
        self.memories["knowledge"][key] = {
            "value": value,
            "updated_at": datetime.now().isoformat()
        }
        self._save()
    
    def recall(self, key: str) -> Optional[str]:
        """回忆知识片段"""
        if key in self.memories["knowledge"]:
            return self.memories["knowledge"][key]["value"]
        return None
    
    def search_knowledge(self, keyword: str) -> list[dict]:
        """搜索知识库"""
        results = []
        for key, entry in self.memories["knowledge"].items():
            if keyword.lower() in key.lower() or \
               keyword.lower() in entry["value"].lower():
                results.append({"key": key, "value": entry["value"]})
        return results

# 在 SimpleAgent 中添加记忆支持

class SimpleAgentWithMemory(SimpleAgent):
    """带记忆的 AI Agent"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.conversation_memory = ConversationMemory()
        self.persistent_memory = PersistentMemory()
        self.session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    def run(self, user_input: str, max_iterations: int = 5) -> str:
        # 使用 conversation_memory 替代 self.messages
        self.conversation_memory.add_message("user", user_input)
        
        for iteration in range(max_iterations):
            response = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_memory.get_context(),
                tools=self.tools,
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            
            if message.tool_calls:
                self.conversation_memory.add_message("assistant", 
                    json.dumps([{"name": tc.function.name, "args": tc.function.arguments} 
                                for tc in message.tool_calls]))
                for tool_call in message.tool_calls:
                    func_name = tool_call.function.name
                    func_args = json.loads(tool_call.function.arguments)
                    result = self._execute_tool(func_name, func_args)
                    self.conversation_memory.add_message("tool", result)
            else:
                final_answer = message.content
                self.conversation_memory.add_message("assistant", final_answer)
                # 保存会话
                self.persistent_memory.save_session(
                    self.session_id, 
                    self.conversation_memory.messages
                )
                return final_answer
        
        return "Agent 已达到最大迭代次数。"

# agent/workflow.py
import time
import json
from typing import Callable, Any
from enum import Enum

class WorkflowStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"

class WorkflowStep:
    """工作流中的一个步骤"""
    
    def __init__(self, name: str, action: Callable, 
                 depends_on: list[str] = None,
                 max_retries: int = 3,
                 timeout: int = 30):
        self.name = name
        self.action = action
        self.depends_on = depends_on or []
        self.max_retries = max_retries
        self.timeout = timeout
        self.status = WorkflowStatus.PENDING
        self.result = None
        self.error = None
        self.retry_count = 0

class WorkflowEngine:
    """多步骤工作流引擎，支持条件分支和错误处理"""
    
    def __init__(self, agent: Any):
        self.agent = agent
        self.steps: dict[str, WorkflowStep] = {}
        self.context: dict = {}  # 工作流上下文，用于步骤间共享数据
    
    def add_step(self, step: WorkflowStep):
        """添加一个工作流步骤"""
        self.steps[step.name] = step
    
    def add_conditional_step(self, name: str, 
                              condition: Callable[[dict], bool],
                              if_true: WorkflowStep,
                              if_false: WorkflowStep = None):
        """添加条件分支步骤"""
        def conditional_action(ctx):
            if condition(ctx):
                return if_true.action(ctx)
            elif if_false:
                return if_false.action(ctx)
            return {"skipped": True, "reason": "Condition not met"}
        
        step = WorkflowStep(
            name=name,
            action=conditional_action,
            depends_on=if_true.depends_on
        )
        self.steps[name] = step
    
    def run(self, initial_input: str) -> dict:
        """执行工作流"""
        results = {}
        self.context["user_input"] = initial_input
        
        # 拓扑排序：按依赖关系排序步骤
        sorted_steps = self._topological_sort()
        
        for step_name in sorted_steps:
            step = self.steps[step_name]
            
            # 检查依赖是否已满足
            if not self._check_dependencies(step):
                step.status = WorkflowStatus.SKIPPED
                results[step_name] = {"status": "skipped", "reason": "依赖未满足"}
                continue
            
            step.status = WorkflowStatus.RUNNING
            print(f"\n▶ 执行步骤: {step.name}")
            
            # 带重试的执行
            for attempt in range(step.max_retries + 1):
                try:
                    start_time = time.time()
                    result = step.action(self.context)
                    elapsed = time.time() - start_time
                    
                    step.status = WorkflowStatus.SUCCESS
                    step.result = result
                    results[step_name] = {
                        "status": "success",
                        "result": result,
                        "elapsed_seconds": round(elapsed, 2)
                    }
                    self.context[step_name] = result
                    print(f"  ✓ 完成 ({elapsed:.1f}s)")
                    break
                    
                except Exception as e:
                    step.retry_count = attempt + 1
                    if attempt < step.max_retries:
                        wait = 2 ** attempt  # 指数退避
                        print(f"  ✗ 失败: {e}，{wait}s 后重试 ({attempt + 1}/{step.max_retries})")
                        time.sleep(wait)
                    else:
                        step.status = WorkflowStatus.FAILED
                        step.error = str(e)
                        results[step_name] = {
                            "status": "failed",
                            "error": str(e)
                        }
                        print(f"  ✗ 最终失败: {e}")
        
        return results
    
    def _check_dependencies(self, step: WorkflowStep) -> bool:
        """检查步骤的依赖是否全部成功"""
        for dep_name in step.depends_on:
            if dep_name not in self.steps:
                return False
            if self.steps[dep_name].status != WorkflowStatus.SUCCESS:
                return False
        return True
    
    def _topological_sort(self) -> list[str]:
        """简单的拓扑排序"""
        visited = set()
        sorted_list = []
        
        def dfs(node):
            if node in visited:
                return
            visited.add(node)
            step = self.steps.get(node)
            if step:
                for dep in step.depends_on:
                    if dep in self.steps:
                        dfs(dep)
                sorted_list.append(node)
        
        for name in self.steps:
            dfs(name)
        
        return sorted_list

# examples/info_gathering_workflow.py
from agent.core import SimpleAgentWithMemory
from agent.workflow import WorkflowEngine, WorkflowStep
from agent.tools import TOOL_REGISTRY

def main():
    agent = SimpleAgentWithMemory()
    engine = WorkflowEngine(agent)
    
    # 步骤 1：搜索收集信息
    def search_info(ctx):
        topic = ctx.get("user_input", "")
        result = agent.run(f"请搜索关于 '{topic}' 的最新信息")
        return result
    
    # 步骤 2：总结信息
    def summarize(ctx):
        raw_info = ctx.get("search_info", "")
        result = agent.run(f"请总结以下信息，提取关键点：\n{raw_info}")
        return result
    
    # 步骤 3：生成报告
    def generate_report(ctx):
        summary = ctx.get("summarize", "")
        result = agent.run(f"请根据以下总结，生成一份结构化的中文报告：\n{summary}")
        return result
    
    # 步骤 4：保存报告（条件分支）
    def save_report(ctx):
        report = ctx.get("generate_report", "")
        with open(f"report_{int(time.time())}.md", "w", encoding="utf-8") as f:
            f.write(report)
        return {"saved": True, "path": f"report_{int(time.time())}.md"}
    
    # 注册步骤
    engine.add_step(WorkflowStep("search_info", search_info))
    engine.add_step(WorkflowStep("summarize", summarize, depends_on=["search_info"]))
    engine.add_step(WorkflowStep("generate_report", generate_report, depends_on=["summarize"]))
    engine.add_step(WorkflowStep("save_report", save_report, depends_on=["generate_report"]))
    
    # 运行工作流
    topic = input("请输入要调研的主题：")
    print(f"\n开始执行信息收集工作流...")
    print("=" * 50)
    
    results = engine.run(topic)
    
    print("\n" + "=" * 50)
    print("工作流执行报告")
    print("=" * 50)
    for step_name, result in results.items():
        status_icon = "✅" if result["status"] == "success" else "❌" if result["status"] == "failed" else "⏭️"
        print(f"{status_icon} {step_name}: {result['status']}")
        if "elapsed_seconds" in result:
            print(f"   耗时: {result['elapsed_seconds']}s")

# 高级重试策略：指数退避 + 抖动
def retry_with_backoff(func, max_retries=3, base_delay=1.0):
    """
    带指数退避和随机抖动的重试机制。
    
    等待时间 = base_delay * (2 ** attempt) + random(0, 1)
    """
    import random
    
    for attempt in range(max_retries + 1):
        try:
            return func()
        except Exception as e:
            if attempt == max_retries:
                raise e
            delay = base_delay * (2 ** attempt) + random.random()
            print(f"  重试 {attempt + 1}/{max_retries}，等待 {delay:.1f}s...")
            time.sleep(delay)
    
    return None

# 使用示例
def safe_web_search(query):
    """带重试的安全搜索"""
    def do_search():
        return TOOL_REGISTRY["web_search"]["function"](query)
    
    try:
        result = retry_with_backoff(do_search, max_retries=3)
        return result
    except Exception as e:
        return {"error": f"搜索最终失败: {str(e)}"}