"""PPO agent wrapper for NephroSim environments.

Uses Stable-Baselines3 PPO with shared hyperparameters from training/configs.py.
This module provides the CLI entrypoint for nephrosim-train.
"""

import sys

import gymnasium as gym
import numpy as np

import nephrosim  # noqa: F401
from nephrosim.training.configs import ENV_CONFIGS
from nephrosim.envs.wrappers import NormalizeObservation, ClipAction


def main():
    """CLI entrypoint: train PPO on a NephroSim environment."""
    try:
        from stable_baselines3 import PPO
    except ImportError:
        print("ERROR: stable-baselines3 required. Install: pip install nephrosim[train]")
        sys.exit(1)

    env_id = sys.argv[1] if len(sys.argv) > 1 else "nephrosim/UreaClearing-v0"

    if env_id not in ENV_CONFIGS:
        print(f"Unknown env: {env_id}. Available: {list(ENV_CONFIGS.keys())}")
        sys.exit(1)

    config = ENV_CONFIGS[env_id]
    print(f"Training PPO on {env_id} for {config['total_timesteps']:,} steps...")

    env = ClipAction(NormalizeObservation(gym.make(env_id)))
    np.random.seed(config["seed"])

    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=config["seed"],
        verbose=1,
    )

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

    print("Training complete.")


if __name__ == "__main__":
    main()
