from typing import List, Dict, Any

class ContextManager:
    def __init__(self, max_tokens: int = 8000, reserve_output_tokens: int = 1500):
        self.max_tokens = max_tokens
        self.reserve_output_tokens = reserve_output_tokens
        self.budget = max_tokens - reserve_output_tokens

    def estimate_tokens(self, text: str) -> int:
        # 简单工程估算：中文约 1.5 chars/token，英文约 4 chars/token
        # 生产环境推荐使用 tiktoken 进行精确计算
        return len(text) // 2 + 1

    def calculate_message_tokens(self, msg: Dict[str, Any]) -> int:
        content = msg.get("content") or ""
        tokens = self.estimate_tokens(str(content))
        if "tool_calls" in msg:
            tokens += self.estimate_tokens(str(msg["tool_calls"]))
        return tokens

    def build_safe_context(
        self,
        system_prompt: str,
        initial_user_goal: str,
        current_state_summary: str,
        raw_history: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        # 1. 组装固定不变的锚定头部
        anchors = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"【任务总目标】: {initial_user_goal}"}
        ]
        if current_state_summary:
            anchors.append({
                "role": "system", 
                "content": f"【当前执行状态进度】:\n{current_state_summary}"
            })

        anchor_cost = sum(self.calculate_message_tokens(m) for m in anchors)
        remaining_budget = self.budget - anchor_cost
        if remaining_budget <= 0:
            raise ValueError("锚定前缀与状态摘要已超出总预算，需压缩系统提示或状态")

        # 2. 将历史消息聚合成原子组（Atomic Groups）
        # 规则：普通消息自成一组；assistant(带tool_calls) 与其后续的全部 tool 消息合并为一组
        atomic_groups: List[List[Dict[str, Any]]] = []
        i = 0
        while i < len(raw_history):
            msg = raw_history[i]
            if msg.get("role") == "assistant" and msg.get("tool_calls"):
                group = [msg]
                expected_call_ids = {c["id"] for c in msg["tool_calls"]}
                i += 1
                while i < len(raw_history) and expected_call_ids:
                    next_msg = raw_history[i]
                    if next_msg.get("role") == "tool":
                        group.append(next_msg)
                        expected_call_ids.discard(next_msg.get("tool_call_id"))
                        i += 1
                    else:
                        break
                atomic_groups.append(group)
            else:
                atomic_groups.append([msg])
                i += 1

        # 3. 从最新消息开始倒序装载，直到耗尽预算
        selected_groups: List[List[Dict[str, Any]]] = []
        accumulated_cost = 0

        for group in reversed(atomic_groups):
            group_cost = sum(self.calculate_message_tokens(m) for m in group)
            if accumulated_cost + group_cost > remaining_budget:
                break
            selected_groups.append(group)
            accumulated_cost += group_cost

        # 4. 正序拼接最终的 Messages
        final_messages = list(anchors)
        for group in reversed(selected_groups):
            final_messages.extend(group)

        return final_messages

from pydantic import BaseModel, Field
from typing import List, Dict, Optional
from enum import Enum

class TaskPhase(str, Enum):
    INITIALIZING = "initializing"
    ANALYZING = "analyzing"
    EXECUTING = "executing"
    VERIFYING = "verifying"
    COMPLETED = "completed"
    FAILED = "failed"

class AgentState(BaseModel):
    task_id: str
    phase: TaskPhase = TaskPhase.INITIALIZING
    goal: str
    completed_steps: List[str] = Field(default_factory=list)
    pending_steps: List[str] = Field(default_factory=list)
    observed_facts: Dict[str, Any] = Field(default_factory=dict)
    retry_counters: Dict[str, int] = Field(default_factory=dict)
    requires_human_approval: bool = False
    last_error: Optional[str] = None

    def advance_step(self, step_name: str, result_summary: str):
        if step_name in self.pending_steps:
            self.pending_steps.remove(step_name)
        self.completed_steps.append(step_name)
        self.observed_facts[step_name] = result_summary