#!/usr/bin/env python
"""UTDG RL Training Script with Auto-Launcher Support.

Trains MaskablePPO agent on the UTDG environment with automatic Godot web build launching.
Supports both native and web-based runtime modes with configurable parameters.
"""

import logging
import os
import sys
import traceback
from typing import Optional

# Ensure local project import works even without pip install -e .
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import hydra
import gymnasium as gym
import numpy as np
from sb3_contrib import MaskablePPO
from sb3_contrib.common.maskable.policies import MaskableMultiInputActorCriticPolicy
from sb3_contrib.common.wrappers import ActionMasker
from wandb.integration.sb3 import WandbCallback
from stable_baselines3.common.vec_env import DummyVecEnv, VecMonitor
from stable_baselines3.common.callbacks import CheckpointCallback
from omegaconf import OmegaConf
import wandb

import utdg_env
from utdg_env.experiment.manager import ExperimentManager
from utdg_env.callbacks.save_model_callback import SaveModelCallback
from utdg_env.callbacks.hf_upload_callback import HFUploadCallback
from utdg_env.callbacks.wandb_callback import WandbMetricsCallback, WandbEpisodeCallback, WandbEvalCallback
from utdg_env.callbacks.curriculum_callback import CurriculumCallback
from utdg_env.utils.hydra_loader import load_config, pretty_print_cfg
from utdg_env.env.registry import make_env
import utdg_env.configs.config_schema

from utdg_env.launcher import GodotWebLauncher, GodotNativeLauncher


# ------------------------------------------------------------
# ActionMask extraction
# ------------------------------------------------------------
def mask_fn(env: gym.Env) -> np.ndarray:
    """Extract action mask from the environment.

    ActionMasker wrapper calls this function to determine which actions
    are valid at the current state. Unwraps the environment to access
    the underlying UntitledTowerDefenseEnv.

    Args:
        env: The wrapped Gymnasium environment.

    Returns:
        Boolean numpy array where True indicates valid actions.
    """
    while hasattr(env, "env"):
        env = env.env
    return env.get_action_mask()


# ------------------------------------------------------------
# Launcher setup (supports all runtime modes)
# ------------------------------------------------------------
def setup_launcher(cfg, logger: logging.Logger):
    """Setup and start Godot launcher for any runtime mode.

    Args:
        cfg: Hydra configuration object containing runtime settings.
        logger: Logger instance for output messages.

    Returns:
        Tuple of (launcher instance, connection_info) or (None, None) if launcher disabled.
    """
    mode = cfg.runtime.mode
    launcher_enabled = cfg.runtime.launcher.get("enabled", False)

    if not launcher_enabled:
        logger.info("Launcher disabled - expecting manual Godot client connection")
        return None, None

    # --------------------------------------------------
    # Native mode: Launch desktop Godot app
    # --------------------------------------------------
    if mode == "native":
        godot_path = cfg.runtime.get("godot_path", "builds/UTDG-macOS.app")
        headless = cfg.runtime.launcher.get("headless", False)

        logger.info("[1/2] Launching native Godot application...")
        launcher = GodotNativeLauncher(
            godot_path=godot_path,
            headless=headless,
        )
        launcher.launch()
        logger.info("  ✓ Godot app launched from %s", godot_path)

        logger.info("[2/2] Waiting for Godot to start WebSocket server on port 9876...")
        return launcher, None

    # --------------------------------------------------
    # Web modes: Launch web build with HTTP server
    # --------------------------------------------------
    elif mode.startswith("web"):
        launcher_mode = cfg.runtime.launcher.get("mode", "local-dev")
        build_dir = cfg.runtime.web.get("path", "builds/web")
        http_port = cfg.runtime.web.get("http_port", 8080)
        ws_port = cfg.runtime.server.port
        headless = cfg.runtime.launcher.get("headless", False)
        open_browser = cfg.runtime.launcher.get("browser", True)

        logger.info("[1/4] Starting Godot web launcher (mode=%s)...", launcher_mode)

        launcher = GodotWebLauncher.from_mode(
            launcher_mode,
            build_dir=build_dir,
            http_port=http_port,
            ws_host="localhost",
            ws_port=ws_port,
            headless=headless,
            open_browser=open_browser,
        )

        http_url, _ = launcher.launch()
        logger.info("  ✓ Web launcher started at %s", http_url)
        logger.info("[2/4] WebSocket server: ws://localhost:%d", ws_port)
        logger.info("[3/4] HTTP server: %s", http_url)

        return launcher, http_url

    else:
        logger.warning("Unknown runtime mode: %s - no launcher configured", mode)
        return None, None


# ------------------------------------------------------------
# Environment creation
# ------------------------------------------------------------
def create_environment(cfg, logger: logging.Logger) -> DummyVecEnv:
    """Create and configure the vectorized masked environment.

    Args:
        cfg: Hydra configuration object.
        logger: Logger instance for output messages.

    Returns:
        Configured DummyVecEnv with ActionMasker wrapper.
    """
    # Create environment using gym.make() with config overrides (Option A)
    # Inside lambda pattern creates fresh environment instance
    def make_env_fn():
        env = make_env(cfg)  # Uses gym.make() with cfg kwargs
        return ActionMasker(env, mask_fn)

    env = DummyVecEnv([make_env_fn])
    env = VecMonitor(env)

    logger.info("  ✓ Environment created")
    logger.info("  Observation space: %s", env.observation_space)
    logger.info("  Action space: %s", env.action_space)

    return env


# ------------------------------------------------------------
# Model creation
# ------------------------------------------------------------
def create_model(env: DummyVecEnv, cfg, logger: logging.Logger, tensorboard_log: Optional[str] = None) -> MaskablePPO:
    """Create and configure the MaskablePPO model.

    Args:
        env: The vectorized environment.
        cfg: Hydra configuration object containing model parameters.
        logger: Logger instance for output messages.
        tensorboard_log: Path for TensorBoard logging (enables metric logging).

    Returns:
        Configured MaskablePPO model ready for training.
    """
    model = MaskablePPO(
        MaskableMultiInputActorCriticPolicy,
        env,
        learning_rate=cfg.model.learning_rate,
        gamma=cfg.model.gamma,
        batch_size=cfg.model.batch_size,
        n_steps=cfg.model.n_steps,
        verbose=cfg.training.verbose,
        device=cfg.training.device,
        tensorboard_log=tensorboard_log,
    )

    logger.info("  ✓ MaskablePPO model created")
    logger.info("  Learning rate: %.4f", cfg.model.learning_rate)
    logger.info("  Gamma: %.2f", cfg.model.gamma)
    logger.info("  Batch size: %d", cfg.model.batch_size)
    if tensorboard_log:
        logger.info("  TensorBoard logging: %s", tensorboard_log)

    return model


# ------------------------------------------------------------
# Cleanup logic
# ------------------------------------------------------------
def cleanup_resources(env: Optional[DummyVecEnv], launcher: Optional[object], logger: logging.Logger):
    """Save the trained model to disk.

    Args:
        model: Trained MaskablePPO model.
        save_dir: Directory to save the model.
        filename: Filename for the saved model.
        logger: Logger instance for output messages.
    """
    logger.info("[cleanup] Shutting down resources...")

    if env is not None:
        try:
            env.close()
            logger.info("  ✓ Environment closed")
        except Exception as e:
            logger.warning("  ⚠ Failed to close environment: %s", e)

    if launcher is not None:
        try:
            launcher.stop()
            logger.info("  ✓ Launcher stopped")
        except Exception as e:
            logger.warning("  ⚠ Failed to stop launcher: %s", e)


# ------------------------------------------------------------
# Main training loop
# ------------------------------------------------------------
def main() -> None:
    """Main training loop with launcher support for web mode.

    Orchestrates the complete training pipeline:
    1. Load configuration
    2. Setup web launcher (if web mode)
    3. Create environment
    4. Initialize model
    5. Train model
    6. Save trained model
    7. Cleanup resources
    """
    # --------------------------------------------------
    # Setup logging
    # --------------------------------------------------
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    logger = logging.getLogger(__name__)

    logger.info("=" * 60)
    logger.info("UTDG Training Pipeline Starting")
    logger.info("=" * 60)

    # ------------------------------------------------------------
    # Load Hydra config
    # ------------------------------------------------------------
    # Extract command-line overrides from sys.argv
    import sys
    overrides = [arg for arg in sys.argv[1:] if "=" in arg]
    cfg = load_config(overrides=overrides)
    print(pretty_print_cfg(cfg))

    # Create ExperimentManager
    manager = ExperimentManager(cfg)

    manager.save_config_snapshot()

    env: Optional[DummyVecEnv] = None
    launcher: Optional[object] = None  # Can be GodotWebLauncher or GodotNativeLauncher

    try:
        # ------------------------------------------------------------
        # Launcher setup (auto-launch Godot if enabled)
        # ------------------------------------------------------------
        launcher, connection_info = setup_launcher(cfg, logger)

        # ------------------------------------------------------------
        # Environment
        # ------------------------------------------------------------
        if launcher is not None:
            logger.info("Creating environment (launcher active)...")
        else:
            logger.info("Creating environment (manual Godot connection expected)...")

        env = create_environment(cfg, logger)

        if launcher is not None:
            if connection_info is not None and hasattr(launcher, 'get_instructions'):
                print(launcher.get_instructions(connection_info))
            logger.info("  ⏳ Waiting for Godot client to connect...")

        # ------------------------------------------------------------
        # Model
        # ------------------------------------------------------------
        logger.info("")
        logger.info("[5/7] Creating MaskablePPO model...")

        # Determine tensorboard log path for W&B sync
        # IMPORTANT: Must be set BEFORE wandb.init() so wandb can find the logs
        tensorboard_log_path = None
        if cfg.callbacks.wandb.enabled:
            tensorboard_log_path = f"{cfg.checkpoint.save_path}/tensorboard_logs"
            logger.info("  TensorBoard logs will be saved to: %s", tensorboard_log_path)

        if cfg.runtime.resume and os.path.exists(cfg.runtime.checkpoint_path):
            logger.info("  ✓ Resume-from-checkpoint")
            model = MaskablePPO.load(cfg.runtime.checkpoint_path, env=env, tensorboard_log=tensorboard_log_path)
        else:
            model = create_model(env, cfg, logger, tensorboard_log=tensorboard_log_path)

        # ------------------------------------------------------------
        # Callbacks
        # ------------------------------------------------------------
        callbacks = []

        # Curriculum learning (adaptive episode horizon)
        if cfg.training.curriculum.enabled:
            curriculum_cb = CurriculumCallback(
                initial_max_steps=cfg.training.curriculum.initial_max_steps,
                step_increase=cfg.training.curriculum.step_increase,
                reward_thresholds=cfg.training.curriculum.reward_thresholds,
                window_size=cfg.training.curriculum.window_size,
                max_steps_cap=cfg.training.curriculum.max_steps_cap,
                verbose=cfg.training.verbose,
            )
            callbacks.append(curriculum_cb)
            logger.info("[Curriculum] Enabled with initial_max_steps=%s",
                       cfg.training.curriculum.initial_max_steps or "from env config")

        # Checkpoints
        if cfg.checkpoint.enabled:
            checkpoint_callback = CheckpointCallback(
                save_freq=cfg.checkpoint.save_freq,
                save_path=cfg.checkpoint.save_path,
                name_prefix=cfg.checkpoint.name_prefix,
                save_replay_buffer=cfg.checkpoint.save_replay_buffer,
                save_vecnormalize=cfg.checkpoint.save_vecnormalize,
            )
            callbacks.append(checkpoint_callback)

        # Wandb logging
        if cfg.callbacks.wandb.enabled:
            # Initialize WandB run first (all callbacks will detect and use it)
            wandb.init(
                project=cfg.callbacks.wandb.project,
                entity=cfg.callbacks.wandb.entity,
                name=cfg.callbacks.wandb.run_name,
                tags=list(cfg.callbacks.wandb.tags),
                mode=cfg.callbacks.wandb.mode,
                save_code=cfg.callbacks.wandb.save_code,
                config=OmegaConf.to_container(cfg, resolve=True),
                sync_tensorboard=False,  # CRITICAL: Disable to allow manual step parameter
                monitor_gym=True,  # Auto-upload gym videos
            )

            # WandB callback for gradients and parameters
            wandb_cb = WandbCallback(
                gradient_save_freq=100,
                model_save_path=f"{cfg.checkpoint.save_path}/wandb_models",
                verbose=2,
                log="all",  # Log gradients and parameters
            )
            callbacks.append(wandb_cb)

            # Custom metrics callback - directly logs SB3 metrics (rollout/*, train/*)
            # This bypasses unreliable TensorBoard sync
            wandb_metrics_cb = WandbMetricsCallback(verbose=1)
            callbacks.append(wandb_metrics_cb)

            # Custom episode callback for reward components
            wandb_ep_cb = WandbEpisodeCallback(
                project=cfg.callbacks.wandb.project,
                entity=cfg.callbacks.wandb.entity,
                run_name=cfg.callbacks.wandb.run_name,
                tags=list(cfg.callbacks.wandb.tags),
                mode=cfg.callbacks.wandb.mode,
                save_code=cfg.callbacks.wandb.save_code,
                config=OmegaConf.to_container(cfg, resolve=True),
                verbose=cfg.training.verbose,
            )
            callbacks.append(wandb_ep_cb)

        # Wandb eval logging
        if cfg.callbacks.wandb.eval_enabled:
            eval_cb = WandbEvalCallback(
                eval_env=env,
                n_eval_episodes=cfg.evaluation.eval_episodes,
                eval_freq=cfg.evaluation.eval_freq,
            )
            callbacks.append(eval_cb)

        # Save final model to experiments directory
        # IMPORTANT: This must run BEFORE HFUploadCallback
        save_model_cb = SaveModelCallback(
                save_path=str(manager.paths.model_file),
                verbose=cfg.training.verbose,
            )
        callbacks.append(save_model_cb)

        # HF upload
        if cfg.callbacks.hf_upload.enabled:
            hf_cb = HFUploadCallback(
                manager=manager,
                push_strategy=cfg.callbacks.hf_upload.push_strategy,  # Optional: pass this too
                verbose=cfg.training.verbose,  # Optional: pass this too
            )
            callbacks.append(hf_cb)

        # Evaluate
        #if cfg.evaluation.eval_freq > 0:
        #    eval_callback = EvalCallback(
        #        env,
        #        eval_freq=cfg.evaluation.eval_freq,
        #        n_eval_episodes=cfg.evaluation.eval_episodes,
        #        best_model_save_path=cfg.checkpoint.save_path,
        #    )
        #    callbacks.append(eval_callback)

        # ------------------------------------------------------------
        # Training
        # ------------------------------------------------------------
        total_timesteps = cfg.training.total_timesteps

        logger.info("")
        logger.info("[6/7] Starting training...")
        logger.info("  Total timesteps: %s", f"{total_timesteps:,}")
        logger.info("=" * 60)

        model.learn(
            total_timesteps=total_timesteps,
            callback=callbacks,
            log_interval=cfg.training.log_interval,
            progress_bar=True,
        )

        logger.info("=" * 60)
        logger.info("  ✓ Training complete")

        # ------------------------------------------------------------
        # Save final model via manager
        # ------------------------------------------------------------
        #logger.info("")
        #logger.info("[7/7] Saving model...")

        #final_model_path = manager.save_model(model)

        #logger.info("  ✓ Saved final model → %s", final_model_path)

        #logger.info("")
        #logger.info("=" * 60)
        #logger.info("Training Pipeline Completed Successfully")
        #logger.info("=" * 60)

    except KeyboardInterrupt:
        logger.info("")
        logger.info("=" * 60)
        logger.info("Training interrupted by user (Ctrl+C)")
        logger.info("=" * 60)

    except Exception as e:
        logger.error("")
        logger.error("=" * 60)
        logger.error("Training failed with error:")
        logger.error("%s", str(e))
        logger.error("=" * 60)
        logger.error("")
        traceback.print_exc()

    finally:
        # ------------------------------------------------------------
        # Cleanup
        # ------------------------------------------------------------
        logger.info("")
        cleanup_resources(env, launcher, logger)
        logger.info("Done.")


if __name__ == "__main__":
    main()
