import json
import time
from dataclasses import dataclass, field
from typing import List, Dict, Any, Callable

@dataclass
class EvalCase:
    case_id: str
    prompt: str
    setup_fixture: Callable[[], None]      # 测试前置环境准备
    assertion_hook: Callable[[], bool]    # 测试后置环境断言

@dataclass
class TrajectoryMetric:
    case_id: str
    passed: bool
    total_steps: int
    total_tokens: int
    elapsed_seconds: float
    error_log: List[str] = field(default_factory=list)

class AgentEvaluationRunner:
    def __init__(self, agent_invoker: Callable[[str], Dict[str, Any]]):
        self.agent_invoker = agent_invoker

    def run_benchmark(self, dataset: List[EvalCase]) -> Dict[str, Any]:
        results: List[TrajectoryMetric] = []
        
        for case in dataset:
            print(f"\n[运行评测用例] {case.case_id}...")
            
            # 1. 重置沙箱初始状态
            case.setup_fixture()
            
            start_time = time.time()
            try:
                # 2. 触发 Agent 执行任务
                trace = self.agent_invoker(case.prompt)
                elapsed = time.time() - start_time
                
                # 3. 执行确定性后置断言（检查物理环境而非模型文本）
                assert_passed = case.assertion_hook()
                
                metric = TrajectoryMetric(
                    case_id=case.case_id,
                    passed=assert_passed,
                    total_steps=trace.get("steps", 0),
                    total_tokens=trace.get("tokens", 0),
                    elapsed_seconds=round(elapsed, 2)
                )
            except Exception as e:
                metric = TrajectoryMetric(
                    case_id=case.case_id,
                    passed=False,
                    total_steps=0,
                    total_tokens=0,
                    elapsed_seconds=round(time.time() - start_time, 2),
                    error_log=[str(e)]
                )
            
            results.append(metric)

        # 4. 汇总基准表现
        total_cases = len(results)
        passed_cases = sum(1 for r in results if r.passed)
        avg_steps = sum(r.total_steps for r in results) / total_cases if total_cases else 0
        
        return {
            "summary": {
                "total": total_cases,
                "passed": passed_cases,
                "pass_rate": f"{(passed_cases / total_cases) * 100:.1f}%",
                "avg_steps_per_task": round(avg_steps, 2)
            },
            "details": [r.__dict__ for r in results]
        }

# --- 模拟评测执行测试 ---
if __name__ == "__main__":
    # 模拟环境状态
    sandbox_db = {"users": {"101": {"status": "LOCKED"}}}

    # 构建 Golden Dataset 用例
    suite = [
        EvalCase(
            case_id="CASE_UNLOCK_USER_101",
            prompt="用户 101 的状态被误锁定了，请帮其执行解锁并记录原因。",
            setup_fixture=lambda: sandbox_db["users"].update({"101": {"status": "LOCKED"}}),
            assertion_hook=lambda: sandbox_db["users"]["101"]["status"] == "ACTIVE"
        )
    ]

    # 模拟被测的 Agent 执行结果
    def mock_agent_call(prompt: str):
        # 模拟模型执行解锁动作
        sandbox_db["users"]["101"]["status"] = "ACTIVE"
        return {"steps": 2, "tokens": 450}

    runner = AgentEvaluationRunner(agent_invoker=mock_agent_call)
    benchmark_report = runner.run_benchmark(suite)
    print("\n--- 自动化评测报告 ---")
    print(json.dumps(benchmark_report, ensure_ascii=False, indent=2))