"""Benchmark runner for NephroSim environments."""

import json


import nephrosim  # noqa: F401
from nephrosim.benchmarks.environments import BENCHMARK_ENVS
from nephrosim.agents.random_agent import evaluate_random
from nephrosim.agents.heuristic import evaluate_heuristic


def run_benchmarks(n_episodes: int = 50, seed: int = 42) -> dict:
    """Run benchmark evaluation across all environments.

    Args:
        n_episodes: Episodes per agent per environment.
        seed: Random seed.

    Returns:
        Dict with results per environment.
    """
    results = {}
    for bench in BENCHMARK_ENVS:
        env_id = bench["env_id"]
        env_name = env_id.split("/")[-1]
        print(f"Benchmarking {env_name}...")

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

        results[env_name] = {
            "env_id": env_id,
            "category": bench["category"],
            "random": {
                "mean_reward": round(random_result["mean_reward"], 2),
                "std_reward": round(random_result["std_reward"], 2),
            },
            "heuristic": {
                "mean_reward": round(heuristic_result["mean_reward"], 2),
                "std_reward": round(heuristic_result["std_reward"], 2),
            },
        }

    return results


def main():
    """CLI entrypoint for nephrosim-benchmark."""
    results = run_benchmarks()
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()
