"""PPO training wrapper using Stable-Baselines3."""

from __future__ import annotations

import json
import os
import sys

import gymnasium as gym
import numpy as np


def train_ppo(
    env_id: str,
    total_timesteps: int = 90000,
    save_dir: str = "checkpoints",
    seed: int = 42,
    verbose: int = 0,
) -> dict:
    """Train a PPO agent on a given environment.

    Args:
        env_id: Gymnasium environment ID.
        total_timesteps: Total training timesteps.
        save_dir: Directory to save trained model.
        seed: Random seed.
        verbose: Verbosity level.

    Returns:
        Dict with training metrics.
    """
    try:
        from stable_baselines3 import PPO
        from stable_baselines3.common.callbacks import BaseCallback
    except ImportError:
        print("stable-baselines3 required. Install with: pip install oncosim[train]")
        sys.exit(1)

    from oncosim.envs.wrappers import FlattenObsWrapper

    env = gym.make(env_id)
    env = FlattenObsWrapper(env)

    class RewardCallback(BaseCallback):
        def __init__(self):
            super().__init__()
            self.episode_rewards: list[float] = []
            self.episode_lengths: list[int] = []
            self._current_reward = 0.0
            self._current_length = 0

        def _on_step(self) -> bool:
            self._current_reward += self.locals["rewards"][0]
            self._current_length += 1
            if self.locals["dones"][0]:
                self.episode_rewards.append(self._current_reward)
                self.episode_lengths.append(self._current_length)
                self._current_reward = 0.0
                self._current_length = 0
            return True

    callback = RewardCallback()

    model = PPO(
        "MlpPolicy",
        env,
        seed=seed,
        verbose=verbose,
        learning_rate=3e-4,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        gamma=0.99,
        device="cpu",
    )

    model.learn(total_timesteps=total_timesteps, callback=callback)

    os.makedirs(save_dir, exist_ok=True)
    model_path = os.path.join(save_dir, f"{env_id.replace('/', '_')}")
    model.save(model_path)

    env.close()

    return {
        "env_id": env_id,
        "total_timesteps": total_timesteps,
        "episode_rewards": callback.episode_rewards,
        "episode_lengths": callback.episode_lengths,
        "mean_reward": float(np.mean(callback.episode_rewards[-50:])) if callback.episode_rewards else 0.0,
        "std_reward": float(np.std(callback.episode_rewards[-50:])) if callback.episode_rewards else 0.0,
        "model_path": model_path,
    }


def evaluate_ppo(
    env_id: str,
    model_path: str,
    n_episodes: int = 100,
    seed: int = 42,
) -> dict[str, float]:
    """Evaluate a trained PPO model."""
    try:
        from stable_baselines3 import PPO
    except ImportError:
        print("stable-baselines3 required.")
        sys.exit(1)

    from oncosim.envs.wrappers import FlattenObsWrapper

    env = gym.make(env_id)
    env = FlattenObsWrapper(env)
    model = PPO.load(model_path)

    rewards = []
    for ep in range(n_episodes):
        obs, _ = env.reset(seed=seed + ep)
        total_reward = 0.0
        done = False
        while not done:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            done = terminated or truncated
        rewards.append(total_reward)

    env.close()

    return {
        "mean_reward": float(np.mean(rewards)),
        "std_reward": float(np.std(rewards)),
        "min_reward": float(np.min(rewards)),
        "max_reward": float(np.max(rewards)),
    }


def main():
    """CLI entry point for training."""
    import oncosim  # noqa: F401 - registers envs

    env_ids = [
        "oncosim/BeamSelection-v0",
        "oncosim/DoseFractionation-v0",
        "oncosim/AdaptiveRT-v0",
    ]

    results = {}
    for env_id in env_ids:
        print(f"Training PPO on {env_id}...")
        result = train_ppo(env_id, total_timesteps=90000)
        results[env_id] = result
        print(f"  Mean reward (last 50 eps): {result['mean_reward']:.3f}")

    os.makedirs("results", exist_ok=True)
    with open("results/training_results.json", "w") as f:
        # Convert to serializable format
        serializable = {}
        for k, v in results.items():
            serializable[k] = {
                "mean_reward": v["mean_reward"],
                "std_reward": v["std_reward"],
                "total_timesteps": v["total_timesteps"],
                "num_episodes": len(v["episode_rewards"]),
                "episode_rewards": v["episode_rewards"],
            }
        json.dump(serializable, f, indent=2)

    print("Results saved to results/training_results.json")


if __name__ == "__main__":
    main()
