"""Benchmark runner for OncoSim environments."""

from __future__ import annotations

import json
import os
import sys

import gymnasium as gym
import numpy as np

import oncosim  # noqa: F401
from oncosim.agents.random_agent import RandomAgent, evaluate_random
from oncosim.benchmarks.environments import BENCHMARK_CONFIGS, ENV_ID_MAP


def run_benchmarks(
    tiers: list[str] | None = None,
    n_episodes: int = 50,
    seed: int = 42,
) -> dict:
    """Run benchmarks across all environments and difficulty tiers.

    Args:
        tiers: List of tier names to run. If None, runs all.
        n_episodes: Episodes per tier.
        seed: Random seed.

    Returns:
        Nested dict: env_name -> tier -> results.
    """
    if tiers is None:
        tiers = ["trivial", "easy", "medium", "hard", "expert"]

    results = {}
    for env_name, tier_configs in BENCHMARK_CONFIGS.items():
        results[env_name] = {}
        for tier in tiers:
            if tier not in tier_configs:
                continue

            config = tier_configs[tier]
            env_id = ENV_ID_MAP[env_name]

            # Create env with tier-specific config
            env = gym.make(
                env_id,
                difficulty=config["difficulty"],
                **config.get("kwargs", {}),
            )

            random_results = evaluate_random(env, n_episodes=n_episodes, seed=seed)
            results[env_name][tier] = {
                "random_mean_reward": random_results["mean_reward"],
                "random_std_reward": random_results["std_reward"],
                "difficulty": config["difficulty"],
                "n_episodes": n_episodes,
            }

            env.close()

    return results


def main():
    """CLI entry point for running benchmarks."""
    results = run_benchmarks()

    os.makedirs("results", exist_ok=True)
    with open("results/benchmark_results.json", "w") as f:
        json.dump(results, f, indent=2)

    # Print summary
    print("\nBenchmark Results")
    print("=" * 70)
    for env_name, tiers in results.items():
        print(f"\n{env_name}:")
        for tier, res in tiers.items():
            print(f"  {tier:8s}: mean={res['random_mean_reward']:8.3f} "
                  f"std={res['random_std_reward']:6.3f}")

    print(f"\nResults saved to results/benchmark_results.json")


if __name__ == "__main__":
    main()
