"""Benchmark runner for AnestheSim environments."""

from __future__ import annotations

import json
import sys

from anesthesim.agents.random_agent import evaluate_random
from anesthesim.agents.heuristic import evaluate_heuristic
from anesthesim.benchmarks.environments import BENCHMARK_ENVS


def run_benchmarks(n_episodes: int = 50, seed: int = 42) -> dict:
    """Run all benchmarks and return results."""
    results = {}

    for bench in BENCHMARK_ENVS:
        env_id = bench["env_id"]
        name = bench["name"]
        print(f"\n=== {name} ===")

        random_result = evaluate_random(env_id, n_episodes=n_episodes, seed=seed)
        heuristic_result = evaluate_heuristic(env_id, n_episodes=n_episodes, seed=seed)

        results[name] = {
            "env_id": env_id,
            "random": {
                "mean_reward": random_result["mean_reward"],
                "std_reward": random_result["std_reward"],
            },
            "heuristic": {
                "mean_reward": heuristic_result["mean_reward"],
                "std_reward": heuristic_result["std_reward"],
            },
        }

        print(f"  Random:    {random_result['mean_reward']:.2f} +/- {random_result['std_reward']:.2f}")
        print(f"  Heuristic: {heuristic_result['mean_reward']:.2f} +/- {heuristic_result['std_reward']:.2f}")

    return results


def main() -> None:
    """CLI entrypoint for running benchmarks."""
    results = run_benchmarks()
    print("\n=== Results JSON ===")
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()
