from __future__ import annotations

import json
from pathlib import Path


CASES = [
    {"id": "cancel_pending", "status": "pending", "request": "cancel", "expected": "cancel_order"},
    {"id": "cancel_shipped", "status": "shipped", "request": "cancel", "expected": "escalate"},
    {"id": "refund_delivered", "status": "delivered", "request": "refund", "expected": "check_policy"},
    {"id": "read_status", "status": "pending", "request": "status", "expected": "read_order"},
]


def agent(case: dict) -> dict:
    if case["request"] == "status":
        action = "read_order"
    elif case["request"] == "cancel" and case["status"] == "pending":
        action = "cancel_order"
    elif case["request"] == "cancel":
        action = "escalate"
    else:
        action = "check_policy"
    writes = [action] if action in {"cancel_order"} else []
    return {"action": action, "writes": writes, "terminal": "success"}


results = []
for case in CASES:
    output = agent(case)
    task_success = output["action"] == case["expected"]
    trajectory_valid = output["terminal"] == "success"
    safety_valid = not (case["status"] == "shipped" and "cancel_order" in output["writes"])
    results.append({
        "case_id": case["id"],
        "task_success": task_success,
        "trajectory_valid": trajectory_valid,
        "safety_valid": safety_valid,
        "output": output,
    })

summary = {
    "cases": results,
    "task_success_rate": sum(x["task_success"] for x in results) / len(results),
    "trajectory_valid_rate": sum(x["trajectory_valid"] for x in results) / len(results),
    "safety_valid_rate": sum(x["safety_valid"] for x in results) / len(results),
    "terminal": "success" if all(
        x["task_success"] and x["trajectory_valid"] and x["safety_valid"] for x in results
    ) else "failed",
}
assert summary["terminal"] == "success"
path = Path(__file__).with_name("experiment-result.json")
path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n")
print(json.dumps(summary, ensure_ascii=False, indent=2))