from anthropic import Anthropic
from dataclasses import dataclass, field
from typing import Optional, Any
import asyncio

@dataclass
class AgentState:
    """
    Agent状态容器

    为什么需要专门的状态类？
    1. 封装所有状态，边界清晰
    2. 便于序列化和持久化
    3. 类型提示让代码更健壮
    """
    goal: str                           # 当前目标
    messages: list[dict] = field(default_factory=list)  # 对话历史
    context: dict[str, Any] = field(default_factory=dict) # 共享上下文
    results: list[dict] = field(default_factory=list)    # 执行结果
    iterations: int = 0                 # 迭代计数器
    max_iterations: int = 10           # 最大迭代次数

class SimpleAgent:
    """
    最小可运行Agent

    设计原则：
    - 简单性：先让它跑起来
    - 可观测性：每个步骤都记录
    - 可停止：明确的终止条件
    """

    def __init__(self, api_key: str, tools: list[dict]):
        self.client = Anthropic(api_key=api_key)
        self.tools = tools

    async def run(self, goal: str) -> dict[str, Any]:
        """
        Agent主循环

        流程：
        1. 初始化状态
        2. 进入执行循环
        3. 每次迭代：LLM推理 → 调用工具 → 更新状态
        4. 达到终止条件时退出
        """
        state = AgentState(goal=goal)
        # 把用户目标作为第一条消息塞进上下文
        state.messages.append({
            "role": "user",
            "content": f"目标：{goal}\n\n请决定下一步行动。"
        })

        while not self._should_terminate(state):
            state.iterations += 1

            # 步骤1：LLM推理，传入完整对话历史
            response = await self._think(state)

            # 步骤2：检查是否需要调用工具
            tool_use_blocks = [b for b in response.content if b.type == "tool_use"]
            if tool_use_blocks:
                # 把 assistant 的响应加入历史
                state.messages.append({"role": "assistant", "content": response.content})
                # 执行所有工具调用，按 Anthropic API 规范构造 tool_result
                tool_results = []
                for tool_use in tool_use_blocks:
                    tool_result = await self._execute_tool(tool_use)
                    state.results.append(tool_result)
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": tool_use.id,
                        "content": str(tool_result),
                    })
                # tool_result 必须以 user 角色回传
                state.messages.append({"role": "user", "content": tool_results})
            else:
                # LLM直接回答，任务完成
                text_blocks = [b for b in response.content if b.type == "text"]
                return {
                    "status": "completed",
                    "goal": goal,
                    "result": text_blocks[0].text if text_blocks else "",
                    "iterations": state.iterations
                }

        return {
            "status": "terminated",
            "goal": goal,
            "reason": "max_iterations_reached",
            "iterations": state.iterations
        }

    def _should_terminate(self, state: AgentState) -> bool:
        """判断是否应该终止"""
        # 达到最大迭代次数
        if state.iterations >= state.max_iterations:
            return True
        # 其他终止条件可以在这里添加
        return False

    async def _think(self, state: AgentState) -> Any:
        """LLM推理：传入完整对话历史，让模型能看到之前的工具调用与结果"""
        response = self.client.messages.create(
            model="claude-opus-4-20250514",
            max_tokens=4096,
            messages=state.messages,
            tools=self.tools
        )
        return response

    async def _execute_tool(self, tool_use: Any) -> dict:
        """执行工具调用，返回原始结果（由调用方按 API 规范包装）"""
        tool_name = tool_use.name
        tool_args = tool_use.input

        # 这里应该调用实际的工具
        # 为了简化，省略具体实现
        return {
            "tool": tool_name,
            "args": tool_args,
            "output": f"Tool {tool_name} executed"
        }

class StateManager:
    """
    状态管理器

    职责：
    1. 压缩对话历史
    2. 提取关键信息
    3. 维护工作上下文
    """

    def __init__(self, max_history: int = 10):
        self.max_history = max_history  # 保留最近N轮对话
        self.summaries = []            # 压缩后的摘要
        self.working_memory = {}       # 当前工作内存

    def add_interaction(self, user_msg: str, assistant_msg: str):
        """添加一轮交互"""
        # 1. 加入历史
        self.summaries.append({
            "user": user_msg,
            "assistant": assistant_msg,
            "timestamp": now()
        })

        # 2. 检查是否需要压缩
        if len(self.summaries) > self.max_history:
            self._compress()

    def _compress(self):
        """
        压缩历史

        压缩策略：
        1. 保留最近N轮完整对话
        2. 更早的对话压缩成摘要
        3. 提取关键实体和决策
        """
        recent = self.summaries[-self.max_history:]
        older = self.summaries[:-self.max_history]

        # 用LLM生成摘要
        summary_prompt = f"""
        请总结以下对话的关键信息：
        {self._format_conversation(older)}

        提取：
        1. 用户的主要目标
        2. 已完成的关键步骤
        3. 当前状态
        4. 重要的中间结果
        """

        summary = self.client.messages.create(
            model="claude-opus-4-20250514",
            messages=[{"role": "user", "content": summary_prompt}]
        )

        # 保存摘要，清空旧历史
        self.working_memory["conversation_summary"] = summary.content[0].text
        self.summaries = recent

    def get_context_for_llm(self) -> str:
        """构建发送给LLM的上下文"""
        parts = []

        # 1. 添加摘要（如果有）
        if "conversation_summary" in self.working_memory:
            parts.append(f"对话摘要：{self.working_memory['conversation_summary']}")

        # 2. 添加最近N轮对话
        for item in self.summaries[-self.max_history:]:
            parts.append(f"用户：{item['user']}")
            parts.append(f"助手：{item['assistant']}")

        return "\n".join(parts)

class SmartContextManager:
    """
    智能上下文管理器

    思路：
    - 按相关性给每条历史打分，而不是按时间顺序硬截断
    - 保留与当前任务最相关的上下文
    """

    def __init__(self, max_tokens: int = 150000):
        self.max_tokens = max_tokens
        self.priority_levels = {
            "critical": ["目标", "关键约束", "核心决策"],
            "important": ["中间结果", "用户偏好", "当前状态"],
            "normal": ["一般对话", "解释说明"],
            "discardable": ["问候", "重复确认"]
        }

    def build_context(self, all_items: list[dict], current_task: str) -> str:
        """
        构建当前任务所需的上下文

        算法：
        1. 为每个上下文项计算相关性分数
        2. 按分数排序
        3. 从最高分开始选取，直到达到token限制
        """
        scored_items = []

        for item in all_items:
            relevance = self._calculate_relevance(item, current_task)
            priority = self._get_priority(item)

            # 综合分数 = 相关性 * 优先级权重
            weight = {"critical": 1.0, "important": 0.7, "normal": 0.4, "discardable": 0.0}
            score = relevance * weight.get(priority, 0.5)

            scored_items.append((score, item))

        # 按分数降序排序
        scored_items.sort(key=lambda x: x[0], reverse=True)

        # 选取项直到达到token限制
        selected = []
        total_tokens = 0

        for score, item in scored_items:
            item_tokens = self._estimate_tokens(item)
            if total_tokens + item_tokens <= self.max_tokens:
                selected.append(item)
                total_tokens += item_tokens

        return self._format_selected(selected)

class ErrorType(Enum):
    """Agent可能遇到的错误类型"""

    # LLM相关错误
    LLM_TIMEOUT = "llm_timeout"           # LLM响应超时
    LLM_RATE_LIMIT = "llm_rate_limit"     # API限流
    LLM_INVALID_RESPONSE = "llm_invalid"  # LLM返回无效响应

    # 工具相关错误
    TOOL_NOT_FOUND = "tool_not_found"     # 工具不存在
    TOOL_EXECUTION_FAILED = "tool_failed" # 工具执行失败
    TOOL_TIMEOUT = "tool_timeout"        # 工具执行超时

    # 状态相关错误
    STATE_CORRUPTED = "state_corrupted"  # 状态损坏
    CONTEXT_OVERFLOW = "context_overflow" # 上下文溢出

    # 业务相关错误
    MAX_ITERATIONS = "max_iterations"    # 达到最大迭代
    USER_CANCELLED = "user_cancelled"    # 用户取消

class ErrorAction(Enum):
    """错误处理动作：决定错误发生后Agent的下一步行为"""

    RETRY = "retry"                       # 重试当前操作
    RETRY_WITH_FALLBACK = "retry_fallback"  # 使用备用方案重试
    FAIL = "fail"                         # 终止任务并上报错误
    RETURN_BEST_RESULT = "return_best"    # 返回当前最佳结果
    ESCALATE = "escalate"                 # 上报给上层处理者

class ErrorHandler:
    """
    错误处理器

    设计原则：
    1. 不同错误类型采用不同的处理策略
    2. 记录错误日志便于调试
    3. 在适当时候回退或重试
    """

    def __init__(self):
        self.error_counts = {}

    async def handle(self, error: Exception, state: AgentState) -> ErrorAction:
        """
        处理错误，返回处理动作
        """
        error_type = self._classify(error)

        # 记录错误
        self._log_error(error_type, error, state)

        # 根据错误类型决定动作
        handlers = {
            ErrorType.LLM_TIMEOUT: self._handle_timeout,
            ErrorType.LLM_RATE_LIMIT: self._handle_rate_limit,
            ErrorType.TOOL_EXECUTION_FAILED: self._handle_tool_failure,
            ErrorType.CONTEXT_OVERFLOW: self._handle_context_overflow,
            ErrorType.MAX_ITERATIONS: self._handle_max_iterations,
        }

        handler = handlers.get(error_type, self._handle_unknown)
        return await handler(error, state)

    async def _handle_timeout(self, error, state):
        """超时处理：等待后重试"""
        state.context["retry_count"] = state.context.get("retry_count", 0) + 1

        if state.context["retry_count"] < 3:
            await asyncio.sleep(2 ** state.context["retry_count"])  # 指数退避
            return ErrorAction.RETRY
        else:
            return ErrorAction.FAIL

    async def _handle_rate_limit(self, error, state):
        """限流处理：等待指定时间"""
        retry_after = getattr(error, "retry_after", 60)
        await asyncio.sleep(retry_after)
        return ErrorAction.RETRY

    async def _handle_tool_failure(self, error, state):
        """工具失败处理：尝试备用方案"""
        tool_name = getattr(error, "tool_name", None)

        # 检查是否有备用工具
        if self._has_fallback(tool_name):
            state.context["using_fallback"] = True
            return ErrorAction.RETRY_WITH_FALLBACK

        return ErrorAction.FAIL

    async def _handle_context_overflow(self, error, state):
        """上下文溢出处理：压缩历史"""
        # 调用状态管理器压缩历史
        state_manager = state.context.get("state_manager")
        if state_manager:
            state_manager._compress()
            return ErrorAction.RETRY
        else:
            return ErrorAction.FAIL

    async def _handle_max_iterations(self, error, state):
        """达到最大迭代：返回当前最佳结果"""
        return ErrorAction.RETURN_BEST_RESULT

class RetryPolicy:
    """
    重试策略

    为什么需要重试策略？
    1. 瞬时故障（网络抖动）可能自行恢复
    2. 限流错误等待后通常可以继续
    3. 合理重试能提高系统稳定性
    """

    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        exponential_base: float = 2.0,
        max_delay: float = 60.0,
        jitter: bool = True
    ):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.exponential_base = exponential_base
        self.max_delay = max_delay
        self.jitter = jitter

    def get_delay(self, attempt: int) -> float:
        """计算重试延迟"""
        # 指数退避
        delay = self.base_delay * (self.exponential_base ** attempt)
        delay = min(delay, self.max_delay)

        # 添加随机抖动，避免惊群效应
        if self.jitter:
            import random
            delay = delay * (0.5 + random.random())

        return delay

class FallbackManager:
    """
    备用方案管理器

    思路：
    - 每个主工具可以有多个备用工具
    - 主工具失败时，自动尝试备用工具
    - 记录使用情况，便于优化
    """

    def __init__(self):
        self.fallback_map = {
            "primary_search": ["fallback_search_1", "fallback_search_2"],
            "primary_translate": ["fallback_translate"],
            "primary_code_exec": ["fallback_sandbox"]
        }
        self.usage_stats = {}

    async def execute_with_fallback(self, primary_tool: str, args: dict) -> Any:
        """执行主工具，失败时尝试备用方案"""
        tools_to_try = [primary_tool] + self.fallback_map.get(primary_tool, [])

        last_error = None
        for tool in tools_to_try:
            try:
                result = await self._execute_tool(tool, args)
                self._record_success(tool)
                return result
            except Exception as e:
                last_error = e
                self._record_failure(tool, e)
                continue

        # 所有工具都失败
        raise AllToolsFailedError(tools_to_try, last_error)

class PermissionScope(Enum):
    """权限范围"""
    NONE = "none"           # 无任何权限
    READ = "read"           # 只读
    WRITE = "write"         # 读写
    EXECUTE = "execute"     # 执行
    ADMIN = "admin"         # 管理

class Permission:
    """
    权限定义

    为什么需要权限系统？
    1. 防止Agent执行危险操作
    2. 限制资源访问范围
    3. 满足合规审计要求
    """

    def __init__(
        self,
        scopes: list[PermissionScope] = None,  # 权限范围列表
        file_paths: list[str] = [],      # 允许访问的文件路径
        allowed_tools: list[str] = [],   # 允许使用的工具
        allowed_domains: list[str] = [], # 允许访问的网络域名
        max_execution_time: int = 300,   # 最大执行时间（秒）
        max_api_calls: int = 100         # 最大API调用次数
    ):
        self.scopes = scopes if scopes is not None else [PermissionScope.READ]
        self.file_paths = file_paths
        self.allowed_tools = allowed_tools
        self.allowed_domains = allowed_domains
        self.max_execution_time = max_execution_time
        self.max_api_calls = max_api_calls

class SecurityManager:
    """
    安全管理器

    职责：
    1. 验证操作权限
    2. 拦截危险操作
    3. 记录所有操作日志
    """

    def __init__(self, permission: Permission):
        self.permission = permission
        self.audit_log = []

    def check_file_access(self, path: str, mode: str) -> bool:
        """检查文件访问权限"""
        # 如果是只读权限但要求写操作
        if mode == "write" and PermissionScope.WRITE not in self.permission.scopes:
            return False

        # 检查路径是否在允许范围内
        import os
        real_path = os.path.realpath(path)

        for allowed_path in self.permission.file_paths:
            if real_path.startswith(os.path.realpath(allowed_path)):
                return True

        return False

    def check_tool_usage(self, tool_name: str) -> bool:
        """检查工具使用权限"""
        return tool_name in self.permission.allowed_tools

    def check_network_access(self, domain: str) -> bool:
        """检查网络访问权限"""
        for allowed in self.permission.allowed_domains:
            if domain.endswith(allowed) or domain == allowed:
                return True
        return False

    def audit(self, operation: str, details: dict):
        """记录审计日志"""
        self.audit_log.append({
            "timestamp": now(),
            "operation": operation,
            "details": details
        })

class SandboxConfig:
    """
    沙箱配置

    为什么需要沙箱？
    即使Agent出错，也只会影响沙箱内的模拟环境
    不会影响真实系统
    """

    def __init__(
        self,
        use_sandbox: bool = True,
        network_isolation: bool = True,
        filesystem_boundary: str = "/workspace/sandbox",
        memory_limit: str = "2GB"
    ):
        self.use_sandbox = use_sandbox
        self.network_isolation = network_isolation
        self.filesystem_boundary = filesystem_boundary
        self.memory_limit = memory_limit

class SandboxExecutor:
    """沙箱执行器"""

    def __init__(self, config: SandboxConfig):
        self.config = config

    async def execute(self, code: str, language: str) -> Any:
        """在沙箱中执行代码"""
        if self.config.use_sandbox:
            return await self._execute_in_sandbox(code, language)
        else:
            return await self._execute_direct(code, language)

    async def _execute_in_sandbox(self, code: str, language: str) -> Any:
        """在沙箱中执行"""
        # 1. 准备沙箱环境
        sandbox = await self._prepare_sandbox()

        # 2. 设置资源限制
        sandbox.set_memory_limit(self.config.memory_limit)
        sandbox.set_network_isolation(self.config.network_isolation)
        sandbox.set_filesystem_boundary(self.config.filesystem_boundary)

        # 3. 执行代码
        try:
            result = await sandbox.run(code, language)
            return result
        finally:
            # 4. 清理沙箱
            await sandbox.cleanup()

import os
from anthropic import Anthropic

client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

tools = [{
    "name": "get_weather",
    "description": "查询指定城市的天气",
    "input_schema": {
        "type": "object",
        "properties": {
            "city": {"type": "string", "description": "城市名"}
        },
        "required": ["city"]
    }
}]

def execute_tool(tool_use):
    """实际工具实现：这里用假数据，生产环境替换成真实 API"""
    if tool_use.name == "get_weather":
        city = tool_use.input["city"]
        return f"{city} 今天晴，25°C"
    raise ValueError(f"未知工具: {tool_use.name}")

def run_agent(goal: str, max_iterations: int = 10):
    messages = [{"role": "user", "content": goal}]

    for i in range(max_iterations):
        response = client.messages.create(
            model="claude-opus-4-20250514",
            max_tokens=1024,
            tools=tools,
            messages=messages,
        )

        tool_use_blocks = [b for b in response.content if b.type == "tool_use"]
        if not tool_use_blocks:
            # 没有工具调用，说明 LLM 给出了最终回答
            text = "".join(b.text for b in response.content if b.type == "text")
            print(f"[最终回答] {text}")
            return text

        # 把 assistant 响应加入历史
        messages.append({"role": "assistant", "content": response.content})

        # 执行所有工具调用，按 API 规范回传 tool_result
        tool_results = []
        for tool_use in tool_use_blocks:
            result = execute_tool(tool_use)
            print(f"[工具调用] {tool_use.name}({tool_use.input}) → {result}")
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": tool_use.id,
                "content": result,
            })
        messages.append({"role": "user", "content": tool_results})

    print("达到最大迭代次数，Agent 终止")

if __name__ == "__main__":
    run_agent("帮我查一下北京和上海的天气，然后总结哪个更适合出行")