Agent/agent/trace/goal_models.py

"""
Goal 数据模型

Goal: 执行计划中的目标节点
GoalTree: 目标树，管理整个执行计划
GoalStats: 目标统计信息
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Dict, Any, List, Optional, Literal
import json


# Goal 状态
GoalStatus = Literal["pending", "in_progress", "completed", "abandoned"]

# Goal 类型
GoalType = Literal["normal", "agent_call"]


@dataclass
class GoalStats:
    """目标统计信息"""
    message_count: int = 0               # 消息数量
    total_tokens: int = 0                # Token 总数
    total_cost: float = 0.0              # 总成本
    preview: Optional[str] = None        # 工具调用摘要，如 "read_file → edit_file → bash"

    def to_dict(self) -> Dict[str, Any]:
        return {
            "message_count": self.message_count,
            "total_tokens": self.total_tokens,
            "total_cost": self.total_cost,
            "preview": self.preview,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "GoalStats":
        return cls(
            message_count=data.get("message_count", 0),
            total_tokens=data.get("total_tokens", 0),
            total_cost=data.get("total_cost", 0.0),
            preview=data.get("preview"),
        )


@dataclass
class Goal:
    """
    执行目标

    使用扁平列表 + parent_id 构建层级结构。
    agent_call 类型用于标记启动了 Sub-Trace 的 Goal。
    """
    id: str                                  # 内部唯一 ID，纯自增（"1", "2", "3"...）
    description: str                         # 目标描述
    reason: str = ""                         # 创建理由（为什么做）
    parent_id: Optional[str] = None          # 父 Goal ID（层级关系）
    type: GoalType = "normal"                # Goal 类型
    status: GoalStatus = "pending"           # 状态
    summary: Optional[str] = None            # 完成/放弃时的总结

    # agent_call 特有
    sub_trace_ids: Optional[List[Dict[str, str]]] = None      # 启动的 Sub-Trace 信息 [{"trace_id": "...", "mission": "..."}]
    agent_call_mode: Optional[str] = None          # "explore" | "delegate" | "sequential"
    sub_trace_metadata: Optional[Dict[str, Dict[str, Any]]] = None  # Sub-Trace 元数据

    # 统计（后端维护，用于可视化边的数据）
    self_stats: GoalStats = field(default_factory=GoalStats)          # 自身统计（仅直接关联的 messages）
    cumulative_stats: GoalStats = field(default_factory=GoalStats)    # 累计统计（自身 + 所有后代）

    created_at: datetime = field(default_factory=datetime.now)

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "id": self.id,
            "description": self.description,
            "reason": self.reason,
            "parent_id": self.parent_id,
            "type": self.type,
            "status": self.status,
            "summary": self.summary,
            "sub_trace_ids": self.sub_trace_ids,
            "agent_call_mode": self.agent_call_mode,
            "sub_trace_metadata": self.sub_trace_metadata,
            "self_stats": self.self_stats.to_dict(),
            "cumulative_stats": self.cumulative_stats.to_dict(),
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Goal":
        """从字典创建"""
        created_at = data.get("created_at")
        if isinstance(created_at, str):
            created_at = datetime.fromisoformat(created_at)

        self_stats = data.get("self_stats", {})
        if isinstance(self_stats, dict):
            self_stats = GoalStats.from_dict(self_stats)

        cumulative_stats = data.get("cumulative_stats", {})
        if isinstance(cumulative_stats, dict):
            cumulative_stats = GoalStats.from_dict(cumulative_stats)

        return cls(
            id=data["id"],
            description=data["description"],
            reason=data.get("reason", ""),
            parent_id=data.get("parent_id"),
            type=data.get("type", "normal"),
            status=data.get("status", "pending"),
            summary=data.get("summary"),
            sub_trace_ids=data.get("sub_trace_ids"),
            agent_call_mode=data.get("agent_call_mode"),
            sub_trace_metadata=data.get("sub_trace_metadata"),
            self_stats=self_stats,
            cumulative_stats=cumulative_stats,
            created_at=created_at or datetime.now(),
        )


@dataclass
class GoalTree:
    """
    目标树 - 管理整个执行计划

    使用扁平列表 + parent_id 构建层级结构
    """
    mission: str                             # 总任务描述
    goals: List[Goal] = field(default_factory=list)  # 扁平列表（通过 parent_id 构建层级）
    current_id: Optional[str] = None         # 当前焦点 goal ID
    _next_id: int = 1                        # 内部 ID 计数器（私有字段）
    created_at: datetime = field(default_factory=datetime.now)

    def find(self, goal_id: str) -> Optional[Goal]:
        """按 ID 查找 Goal"""
        for goal in self.goals:
            if goal.id == goal_id:
                return goal
        return None

    def find_by_display_id(self, display_id: str) -> Optional[Goal]:
        """按显示 ID 查找 Goal（如 "1", "2.1", "2.2"）"""
        for goal in self.goals:
            if self._generate_display_id(goal) == display_id:
                return goal
        return None

    def find_parent(self, goal_id: str) -> Optional[Goal]:
        """查找指定 Goal 的父节点"""
        goal = self.find(goal_id)
        if not goal or not goal.parent_id:
            return None
        return self.find(goal.parent_id)

    def get_children(self, parent_id: Optional[str]) -> List[Goal]:
        """获取指定父节点的所有子节点"""
        return [g for g in self.goals if g.parent_id == parent_id]

    def get_current(self) -> Optional[Goal]:
        """获取当前焦点 Goal"""
        if self.current_id:
            return self.find(self.current_id)
        return None

    def _generate_id(self) -> str:
        """生成新的 Goal ID（纯自增）"""
        new_id = str(self._next_id)
        self._next_id += 1
        return new_id

    def _generate_display_id(self, goal: Goal) -> str:
        """生成显示序号（1, 2, 2.1, 2.2...）"""
        if not goal.parent_id:
            # 顶层目标：找到在同级中的序号
            siblings = [g for g in self.goals if g.parent_id is None and g.status != "abandoned"]
            try:
                index = [g.id for g in siblings].index(goal.id) + 1
                return str(index)
            except ValueError:
                return "?"
        else:
            # 子目标：父序号 + "." + 在同级中的序号
            parent = self.find(goal.parent_id)
            if not parent:
                return "?"
            parent_display = self._generate_display_id(parent)
            siblings = [g for g in self.goals if g.parent_id == goal.parent_id and g.status != "abandoned"]
            try:
                index = [g.id for g in siblings].index(goal.id) + 1
                return f"{parent_display}.{index}"
            except ValueError:
                return f"{parent_display}.?"

    def add_goals(
        self,
        descriptions: List[str],
        reasons: Optional[List[str]] = None,
        parent_id: Optional[str] = None
    ) -> List[Goal]:
        """
        添加目标

        如果 parent_id 为 None，添加到顶层
        如果 parent_id 有值，添加为该 goal 的子目标
        """
        if parent_id:
            parent = self.find(parent_id)
            if not parent:
                raise ValueError(f"Parent goal not found: {parent_id}")

        # 创建新目标
        new_goals = []
        for i, desc in enumerate(descriptions):
            goal_id = self._generate_id()
            reason = reasons[i] if reasons and i < len(reasons) else ""
            goal = Goal(
                id=goal_id,
                description=desc.strip(),
                reason=reason,
                parent_id=parent_id
            )
            self.goals.append(goal)
            new_goals.append(goal)

        return new_goals

    def add_goals_after(
        self,
        target_id: str,
        descriptions: List[str],
        reasons: Optional[List[str]] = None
    ) -> List[Goal]:
        """
        在指定 Goal 后面添加兄弟节点

        新创建的 goals 与 target 有相同的 parent_id，
        并插入到 goals 列表中 target 的后面。
        """
        target = self.find(target_id)
        if not target:
            raise ValueError(f"Target goal not found: {target_id}")

        # 创建新 goals（parent_id 与 target 相同）
        new_goals = []
        for i, desc in enumerate(descriptions):
            goal_id = self._generate_id()
            reason = reasons[i] if reasons and i < len(reasons) else ""
            goal = Goal(
                id=goal_id,
                description=desc.strip(),
                reason=reason,
                parent_id=target.parent_id  # 同层级
            )
            new_goals.append(goal)

        # 插入到 target 后面（调整 goals 列表顺序）
        target_index = self.goals.index(target)
        for i, goal in enumerate(new_goals):
            self.goals.insert(target_index + 1 + i, goal)

        return new_goals

    def focus(self, goal_id: str) -> Goal:
        """切换焦点到指定 Goal，并将其状态设为 in_progress"""
        goal = self.find(goal_id)
        if not goal:
            raise ValueError(f"Goal not found: {goal_id}")

        # 更新状态
        if goal.status == "pending":
            goal.status = "in_progress"

        self.current_id = goal_id
        return goal

    def complete(self, goal_id: str, summary: str, clear_focus: bool = True) -> Goal:
        """
        完成指定 Goal

        Args:
            goal_id: 要完成的目标 ID
            summary: 完成总结
            clear_focus: 如果完成的是当前焦点，是否清除焦点（默认 True）
        """
        goal = self.find(goal_id)
        if not goal:
            raise ValueError(f"Goal not found: {goal_id}")

        goal.status = "completed"
        goal.summary = summary

        # 如果完成的是当前焦点，根据参数决定是否清除焦点
        if clear_focus and self.current_id == goal_id:
            self.current_id = None

        # 检查是否所有兄弟都完成了，如果是则自动完成父节点
        if goal.parent_id:
            siblings = self.get_children(goal.parent_id)
            all_completed = all(g.status == "completed" for g in siblings)
            if all_completed:
                parent = self.find(goal.parent_id)
                if parent and parent.status != "completed":
                    # 自动级联完成父节点
                    parent.status = "completed"
                    if not parent.summary:
                        parent.summary = "所有子目标已完成"

        return goal

    def abandon(self, goal_id: str, reason: str) -> Goal:
        """放弃指定 Goal"""
        goal = self.find(goal_id)
        if not goal:
            raise ValueError(f"Goal not found: {goal_id}")

        goal.status = "abandoned"
        goal.summary = reason

        # 如果放弃的是当前焦点，清除焦点
        if self.current_id == goal_id:
            self.current_id = None

        return goal

    def to_prompt(self, include_abandoned: bool = False, include_summary: bool = False) -> str:
        """
        格式化为 Prompt 注入文本

        Args:
            include_abandoned: 是否包含已废弃的目标
            include_summary: 是否显示 completed/abandoned goals 的 summary 详情
                False（默认）= 精简视图，用于日常周期性注入
                True = 完整视图（含 summary），用于压缩时提供上下文

        展示策略：
        - 过滤掉 abandoned 目标（除非明确要求）
        - 完整展示所有顶层目标
        - 完整展示当前 focus 目标的父链及其所有子孙
        - 其他分支的子目标折叠显示（只显示数量和状态）
        - include_summary=True 时不折叠，全部展开并显示 summary
        """
        lines = []
        lines.append(f"**Mission**: {self.mission}")

        if self.current_id:
            current = self.find(self.current_id)
            if current:
                display_id = self._generate_display_id(current)
                lines.append(f"**Current**: {display_id} {current.description}")

        lines.append("")
        lines.append("**Progress**:")

        # 获取当前焦点的祖先链（从根到当前节点的路径）
        current_path = set()
        if self.current_id:
            goal = self.find(self.current_id)
            while goal:
                current_path.add(goal.id)
                if goal.parent_id:
                    goal = self.find(goal.parent_id)
                else:
                    break

        def format_goal(goal: Goal, indent: int = 0) -> List[str]:
            # 跳过废弃的目标（除非明确要求包含）
            if goal.status == "abandoned" and not include_abandoned:
                return []

            prefix = "    " * indent

            # 状态图标
            if goal.status == "completed":
                icon = "[✓]"
            elif goal.status == "in_progress":
                icon = "[→]"
            elif goal.status == "abandoned":
                icon = "[✗]"
            else:
                icon = "[ ]"

            # 生成显示序号
            display_id = self._generate_display_id(goal)

            # 当前焦点标记
            current_mark = " ← current" if goal.id == self.current_id else ""

            result = [f"{prefix}{icon} {display_id}. {goal.description}{current_mark}"]

            # 显示 summary：include_summary=True 时全部显示，否则只在焦点路径上显示
            if goal.summary and (include_summary or goal.id in current_path):
                result.append(f"{prefix}    → {goal.summary}")

            # 递归处理子目标
            children = self.get_children(goal.id)

            # include_summary 模式下不折叠，全部展开
            if include_summary:
                for child in children:
                    result.extend(format_goal(child, indent + 1))
                return result

            # 判断是否需要折叠
            # 如果当前 goal 或其子孙在焦点路径上，完整展示
            should_expand = goal.id in current_path or any(
                child.id in current_path for child in self._get_all_descendants(goal.id)
            )

            if should_expand or not children:
                # 完整展示子目标
                for child in children:
                    result.extend(format_goal(child, indent + 1))
            else:
                # 折叠显示：只显示子目标的统计
                non_abandoned = [c for c in children if c.status != "abandoned"]
                if non_abandoned:
                    completed = sum(1 for c in non_abandoned if c.status == "completed")
                    in_progress = sum(1 for c in non_abandoned if c.status == "in_progress")
                    pending = sum(1 for c in non_abandoned if c.status == "pending")

                    status_parts = []
                    if completed > 0:
                        status_parts.append(f"{completed} completed")
                    if in_progress > 0:
                        status_parts.append(f"{in_progress} in progress")
                    if pending > 0:
                        status_parts.append(f"{pending} pending")

                    status_str = ", ".join(status_parts)
                    result.append(f"{prefix}    ({len(non_abandoned)} subtasks: {status_str})")

            return result

        # 处理所有顶层目标
        top_goals = self.get_children(None)
        for goal in top_goals:
            lines.extend(format_goal(goal))

        return "\n".join(lines)

    def _get_all_descendants(self, goal_id: str) -> List[Goal]:
        """获取指定 Goal 的所有子孙节点"""
        descendants = []
        children = self.get_children(goal_id)
        for child in children:
            descendants.append(child)
            descendants.extend(self._get_all_descendants(child.id))
        return descendants

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "mission": self.mission,
            "goals": [g.to_dict() for g in self.goals],
            "current_id": self.current_id,
            "_next_id": self._next_id,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "GoalTree":
        """从字典创建"""
        goals = [Goal.from_dict(g) for g in data.get("goals", [])]
        created_at = data.get("created_at")
        if isinstance(created_at, str):
            created_at = datetime.fromisoformat(created_at)

        return cls(
            mission=data["mission"],
            goals=goals,
            current_id=data.get("current_id"),
            _next_id=data.get("_next_id", 1),
            created_at=created_at or datetime.now(),
        )

    def rebuild_for_rewind(self, cutoff_time: datetime) -> "GoalTree":
        """
        为 Rewind 重建干净的 GoalTree

        以截断点消息的 created_at 为界：
        - 保留 created_at <= cutoff_time 的所有 goals（无论状态）
        - 丢弃 cutoff_time 之后创建的 goals
        - 将被保留的 in_progress goals 重置为 pending
        - 清空 current_id，让 Agent 重新选择焦点

        Args:
            cutoff_time: 截断点消息的创建时间

        Returns:
            新的干净 GoalTree
        """
        surviving_goals = []
        for goal in self.goals:
            if goal.created_at <= cutoff_time:
                surviving_goals.append(goal)

        # 清理 parent_id 引用：如果 parent 不在存活列表中，设为 None
        surviving_ids = {g.id for g in surviving_goals}
        for goal in surviving_goals:
            if goal.parent_id and goal.parent_id not in surviving_ids:
                goal.parent_id = None
            # 将 in_progress 重置为 pending（回溯后需要重新执行）
            if goal.status == "in_progress":
                goal.status = "pending"

        new_tree = GoalTree(
            mission=self.mission,
            goals=surviving_goals,
            current_id=None,
            _next_id=self._next_id,
            created_at=self.created_at,
        )
        return new_tree

    def save(self, path: str) -> None:
        """保存到 JSON 文件"""
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)

    @classmethod
    def load(cls, path: str) -> "GoalTree":
        """从 JSON 文件加载"""
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return cls.from_dict(data)