"""Trajectory Recorder — 采集五智能体每次调用的完整轨迹。

轨迹格式 (trajectory):
{
  "trajectory_id": "traj_<timestamp>",
  "customer": "客户名",
  "workflow_id": "...",
  "steps": [
     {
       "agent": "data_curator",
       "input": "任务描述",
       "prompt_snippet": "...",
       "tools": ["read_xingbang_data", ...],
       "output_file": "...",
       "output_preview": "...",
       "success": true/false,
       "time_ms": 1234
     }, ...
  ],
  "final_result": {...},
  "cumulative_reward": 0.0
}

轨迹保存在 <base>/xingbang_evolution/trajectories/*.json。
"""

from __future__ import annotations

import json
import time
import uuid
from pathlib import Path
from typing import Any


class TrajectoryRecorder:
    """记录并持久化一次评审流水线的完整执行轨迹。"""

    def __init__(self, base_dir: Path, enabled: bool = True):
        self._base_dir = Path(base_dir)
        self._dir = self._base_dir / "xingbang_evolution" / "trajectories"
        self._dir.mkdir(parents=True, exist_ok=True)
        self.enabled = enabled
        # 当前正在记录的轨迹
        self._current: dict[str, Any] | None = None
        # metadata (customer 等)
        self._meta: dict = {}
        self._start_time: float = 0.0

    def begin(self, workflow_id: str | None = None, customer: str = "") -> str:
        """开始记录一条新轨迹,返回 trajectory_id。"""
        self._start_time = time.time()
        self._current = {
            "trajectory_id": f"traj_{int(time.time())}_{uuid.uuid4().hex[:6]}",
            "workflow_id": workflow_id or f"wf_{int(time.time())}",
            "customer": customer or "",
            "created_at": self._start_time,
            "steps": [],
            "final_result": {},
            "cumulative_reward": 0.0,
        }
        return self._current["trajectory_id"]

    def add_step(self, step: dict[str, Any]) -> None:
        """追加一步智能体执行记录。"""
        if self._current is not None:
            self._current["steps"].append(step)

    def finalize(self, success: bool = True, summary: str = "") -> dict[str, Any]:
        """结束记录,保存轨迹到磁盘,返回轨迹。"""
        if self._current is None:
            return {}
        self._current["final_result"] = {
            "success": success,
            "summary": summary,
            "duration_ms": int((time.time() - self._start_time) * 1000),
            "total_steps": len(self._current["steps"]),
        }
        # 记录成功/失败
        trajectory = self._current
        self._save(trajectory)
        self._current = None
        return trajectory

    def _save(self, trajectory: dict[str, Any]) -> Path:
        path = self._dir / f"{trajectory['trajectory_id']}.json"
        path.write_text(
            json.dumps(trajectory, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        return path

    # ── 从调用侧采集 ────────────────────────────────────
    def on_evidence(self, kind: str, data: dict) -> None:
        """订阅 agents 的 evidence 回调,累积成本步记录。"""
        if not self.enabled or self._current is None:
            return
        if kind == "agent_start":
            buffered = []
            self._current["_buffered"] = buffered
        # 这一步先不展开, 由 workflow 显式调用 add_step 更为清晰

    def list_trajectories(self) -> list[Path]:
        return sorted(self._dir.glob("*.json"))

    def load(self, trajectory_id: str) -> dict[str, Any]:
        path = self._dir / f"{trajectory_id}.json"
        if path.exists():
            return json.loads(path.read_text(encoding="utf-8"))
        return {}