#!/usr/bin/env python3
"""
Inference rollout for UTDG using a trained MaskablePPO model.

Loads:
    models/model_policy.zip

Runs the agent inside:
    DummyVecEnv + ActionMasker

and steps through live Godot gameplay via WebSockets.

Press CTRL+C to interrupt.
"""

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

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import gymnasium as gym
import numpy as np
from sb3_contrib import MaskablePPO
from sb3_contrib.common.wrappers import ActionMasker
from stable_baselines3.common.vec_env import DummyVecEnv

import utdg_env
from utdg_env.launcher import GodotWebLauncher
from utdg_env.utils.hydra_loader import load_config, pretty_print_cfg


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()


def setup_web_launcher(
    cfg,
    logger: logging.Logger,
) -> tuple[Optional[GodotWebLauncher], Optional[str]]:
    """Setup and start Godot web launcher for web runtime mode.

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

    Returns:
        Tuple of (launcher instance, http_url) or (None, None) if not in web mode.
    """
    if cfg.runtime.mode != "web":
        return None, None

    ws_port = cfg.runtime.server.port
    build_dir = cfg.runtime.get("godot_path", "builds/web")
    headless = cfg.runtime.launcher.get("headless", False)

    logger.info("[1/4] Starting HTTP server for Godot web build...")

    launcher = GodotWebLauncher(
        build_dir=build_dir,
        http_port=8080,
        ws_host="localhost",
        ws_port=ws_port,
        headless=headless,
    )

    http_url, _ = launcher.start_http_server()
    logger.info("  ✓ HTTP server started at %s", http_url)

    if not headless:
        logger.info("[2/4] Launching browser...")
        launcher.launch_browser(http_url)
        logger.info("  ✓ Browser launched")
    else:
        logger.info("[2/4] Skipping browser launch (headless mode)")

    logger.info("[3/4] WebSocket server will start with environment...")
    logger.info("  WebSocket: ws://localhost:%d", ws_port)
    logger.info("  HTTP: %s", http_url)

    return launcher, http_url


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.
    """
    env = DummyVecEnv([
        lambda: ActionMasker(gym.make("UTDGEnv-v0", sb3_mode=False), mask_fn)
    ])

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

    return env


def cleanup_resources(
    env: Optional[DummyVecEnv],
    launcher: Optional[GodotWebLauncher],
    logger: logging.Logger,
) -> None:
    """Clean up environment and launcher resources.

    Args:
        env: The vectorized environment (may be None).
        launcher: The Godot web launcher (may be None).
        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)


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

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

    # --------------------------------------------------
    # [1/6] Load configuration
    # --------------------------------------------------
    logger.info("=" * 60)
    logger.info("UTDG Rollout Pipeline Starting")
    logger.info("=" * 60)

    cfg = load_config()
    print(pretty_print_cfg(cfg))
    print("===================================\n")

    # Initialize resources
    env: Optional[DummyVecEnv] = None
    launcher: Optional[GodotWebLauncher] = None

    try:
        # --------------------------------------------------
        # [2/6] Setup web launcher (if web mode)
        # --------------------------------------------------
        launcher, http_url = setup_web_launcher(cfg, logger)

        # --------------------------------------------------
        # [3/6] Create environment
        # --------------------------------------------------
        if cfg.runtime.mode == "web":
            logger.info("[4/4] Creating environment...")
        else:
            logger.info("[1/1] Creating environment (native mode)...")

        env = create_environment(cfg, logger)

        # Display connection instructions for web mode
        if launcher is not None and http_url is not None:
            print(launcher.get_instructions(http_url))
            logger.info("  ⏳ Waiting for Godot client to connect...")

        # --------------------------------------------------
        # [4/6] Load trained MaskablePPO policy
        # --------------------------------------------------
        logger.info("")
        logger.info("[5/6] Loading trained model...")

        model_path = os.path.join("models", "model_policy.zip")

        if not os.path.exists(model_path):
            raise FileNotFoundError(
                f"Model not found: {model_path}\n"
                f"Make sure training was completed and saved."
            )

        model = MaskablePPO.load(model_path, env=env, device=cfg.model.device)
        logger.info("  ✓ Model loaded from %s", model_path)
        logger.info("  Device: %s", cfg.model.device)

        # --------------------------------------------------
        # [5/6] Reset environment and run inference
        # --------------------------------------------------
        logger.info("")
        logger.info("[6/6] Starting inference...")
        logger.info("=" * 60)

        obs = env.reset()
        logger.info("Environment reset. Beginning agent play...")

        step = 0
        total_reward = 0.0

        while True:
            # Model's predict call accepts masks internally (sb3-contrib feature)
            action, _ = model.predict(obs, deterministic=True)

            # Step the env
            obs, reward, done, info = env.step(action)
            total_reward += reward[0]

            logger.info(
                "[step %d] action=%d reward=%.2f done=%s",
                step, action[0], reward[0], done[0]
            )

            step += 1

            if done[0]:
                logger.info("")
                logger.info("=" * 60)
                logger.info("Episode finished!")
                logger.info("  Total steps: %d", step)
                logger.info("  Total reward: %.2f", total_reward)
                logger.info("=" * 60)
                break

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

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

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


if __name__ == "__main__":
    main()
