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)
        
        while not self._should_terminate(state):
            state.iterations += 1
            
            # 步骤1：LLM推理
            response = await self._think(state)
            
            # 步骤2：检查是否需要调用工具
            if response.content[0].type == "tool_use":
                tool_result = await self._execute_tool(response)
                state.results.append(tool_result)
                state.messages.append(response)
                state.messages.append(tool_result)
            else:
                # LLM直接回答，任务完成
                return {
                    "status": "completed",
                    "goal": goal,
                    "result": response.content[0].text,
                    "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-20241120",
            max_tokens=4096,
            messages=[{
                "role": "user",
                "content": f"目标：{state.goal}\n\n请决定下一步行动。"
            }],
            tools=self.tools
        )
        return response
    
    async def _execute_tool(self, response: Any) -> dict:
        """执行工具调用"""
        tool_use = response.content[0]
        tool_name = tool_use.name
        tool_args = tool_use.input
        
        # 这里应该调用实际的工具
        # 为了简化，省略具体实现
        return {
            "role": "user",
            "content": f"Tool {tool_name} executed with args {tool_args}"
        }

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-20241120",
            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 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,
        file_paths: list[str] = [],      # 允许访问的文件路径
        allowed_tools: list[str] = [],   # 允许使用的工具
        allowed_domains: list[str] = [], # 允许访问的网络域名
        max_execution_time: int = 300,   # 最大执行时间（秒）
        max_api_calls: int = 100         # 最大API调用次数
    ):
        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隔离(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()