"""Minimal training launcher for quick experiments.

Usage:
    python ml/run_train.py --total-timesteps 1500000 --gui
"""

import argparse
import os
import re
import sys
import warnings
from pathlib import Path
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback

sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv
from ml.callbacks import CurriculumCallback

# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")


def get_next_run_number(save_dir: str) -> int:
    """Scans the save directory for existing ppo<number> patterns and returns the next integer."""
    if not os.path.exists(save_dir):
        return 1
    
    existing_numbers = []
    for item in os.listdir(save_dir):
        # Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12)
        matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE)
        for m in matches:
            existing_numbers.append(int(m))
            
    return max(existing_numbers, default=0) + 1


def make_env(rank: int, use_gui: bool = False, seed: int = 0):
    """Utility helper to instantiate parallel JackBot environments."""
    def _init():
        env = JackBotEnv(
            use_gui=use_gui,
            random_command=True,
        )
        env.reset(seed=seed + rank)
        return env
    return _init


def main():
    parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer")
    parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments")
    parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps")
    parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs")
    parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
    parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
    args = parser.parse_args()

    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(args.save_dir, exist_ok=True)

    # Automatically determine next run number (e.g. ppo1, ppo2, ppo3...)
    run_num = get_next_run_number(args.save_dir)
    ppo_name = f"ppo{run_num}"

    print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...")
    
    env_fns = [
        make_env(rank=i, use_gui=(args.gui if i == 0 else False)) 
        for i in range(args.num_workers)
    ]
    vec_env = SubprocVecEnv(env_fns)

    # Initialize PPO Policy Hyperparameters
    model = PPO(
        policy="MlpPolicy",
        env=vec_env,
        learning_rate=1e-4,
        n_steps=256,
        batch_size=256,
        n_epochs=10,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        ent_coef=0.03,
        target_kl=0.05,
        vf_coef=0.5,
        max_grad_norm=0.5,
        verbose=1,
        tensorboard_log=args.log_dir,
        device="cpu", 
    )

    # Setup Callbacks with ppo<number> naming
    checkpoint_callback = CheckpointCallback(
        save_freq=max(1, args.save_freq // args.num_workers),
        save_path=args.save_dir,
        name_prefix=f"jackbot_{ppo_name}",
    )
    curriculum_callback = CurriculumCallback()

    eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
    best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}")

    eval_callback = EvalCallback(
        eval_env,
        best_model_save_path=best_model_path,
        log_path="ml/logs/results",
        eval_freq=max(1, 10_000 // args.num_workers),
        deterministic=True,
        render=False,
    )

    print(f"[Train] Starting training for {args.total_timesteps} timesteps...")
    try:
        model.learn(
            total_timesteps=args.total_timesteps,
            callback=[checkpoint_callback, curriculum_callback, eval_callback],
            progress_bar=True,
        )
        final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")
        model.save(final_model_path)
        print(f"[Train] Training complete! Saved final model to {final_model_path}")
    finally:
        vec_env.close()
        eval_env.close()


if __name__ == "__main__":
    main()