"""PPO agent for AnestheSim environments (CLI entrypoint).

Uses shared hyperparameters from training/configs.py.
"""

from __future__ import annotations

import argparse
import sys


def main() -> None:
    """CLI entrypoint for training PPO on AnestheSim environments."""
    parser = argparse.ArgumentParser(description="Train PPO on AnestheSim environments")
    parser.add_argument(
        "--env",
        type=str,
        default="anesthesim/PropofolControl-v0",
        help="Environment ID to train on",
    )
    parser.add_argument("--output", type=str, default=None, help="Output model path")
    args = parser.parse_args()

    try:
        from stable_baselines3 import PPO
    except ImportError:
        print("stable-baselines3 required. Install: pip install anesthesim[train]")
        sys.exit(1)

    import anesthesim  # noqa: F401
    import gymnasium as gym
    import numpy as np
    import random

    from anesthesim.training.configs import ENV_CONFIGS
    from anesthesim.envs.wrappers import NormalizeObservation, ClipAction

    config = ENV_CONFIGS.get(args.env)
    if config is None:
        print(f"No config for {args.env}. Available: {list(ENV_CONFIGS.keys())}")
        sys.exit(1)

    seed = config["seed"]
    random.seed(seed)
    np.random.seed(seed)
    print(f"Training with seed={seed}")

    env = ClipAction(NormalizeObservation(gym.make(args.env)))

    model = PPO(
        "MlpPolicy",
        env,
        learning_rate=config["learning_rate"],
        n_steps=config["n_steps"],
        batch_size=config["batch_size"],
        gamma=config["gamma"],
        gae_lambda=config["gae_lambda"],
        clip_range=config["clip_range"],
        ent_coef=config["ent_coef"],
        seed=seed,
        verbose=1,
    )

    model.learn(total_timesteps=config["total_timesteps"])

    output_path = args.output or f"results/models/{args.env.split('/')[-1]}_ppo"
    model.save(output_path)
    print(f"Model saved to {output_path}")
    env.close()


if __name__ == "__main__":
    main()
