"""
Script to train an agent using BC (Behavioral Cloning).

Behavioral Cloning directly learns to mimic expert demonstrations through
supervised learning, without adversarial training or reinforcement learning.

How to run:
    python train_bc.py

Requirements:
    pip install imitation
"""
import logging
import os
import sys

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

import gymnasium as gym
import numpy as np
from gymnasium.wrappers import FlattenObservation
from omegaconf import DictConfig, OmegaConf
from stable_baselines3.common.vec_env import DummyVecEnv, VecMonitor
from stable_baselines3.common.evaluation import evaluate_policy
from imitation.algorithms import bc
from imitation.data import rollout, serialize

from utdg_env.utils.hydra_loader import load_config, pretty_print_cfg
from utdg_env.env.registry import make_env
from utdg_env.launcher import GodotWebLauncher, GodotNativeLauncher

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)


def setup_launcher(cfg, logger: logging.Logger):
    """Setup and start Godot launcher for any runtime mode.

    Copied from trainer.py to maintain consistency.
    """
    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

    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)
        return launcher, None

    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)
        return launcher, http_url

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


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

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

    Returns:
        Configured DummyVecEnv with flattened observations.
    """
    def make_env_fn():
        env = make_env(cfg)
        # Flatten Dict observations to a single Box space
        # This is required for compatibility with imitation library
        env = FlattenObservation(env)
        return env

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

    logger.info("  ✓ Environment created (with FlattenObservation wrapper)")
    logger.info("  Observation space: %s", venv.observation_space)
    logger.info("  Action space: %s", venv.action_space)

    return venv


def cleanup_resources(venv, launcher, logger: logging.Logger):
    """Cleanup environment and launcher resources."""
    logger.info("[cleanup] Shutting down resources...")

    if venv is not None:
        try:
            venv.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 train_bc() -> None:
    """Train an agent using Behavioral Cloning."""

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

    # Load config using the same method as trainer.py
    overrides = [arg for arg in sys.argv[1:] if "=" in arg]
    cfg = load_config(overrides=overrides)
    print(pretty_print_cfg(cfg))

    venv = None
    launcher = None

    try:
        # ------------------------------------------------------------
        # Load Demonstrations
        # ------------------------------------------------------------
        trajectory_path = "human_demonstrations.pkl"
        logger.info(f"\n[1/6] Loading demonstrations from {trajectory_path}...")

        try:
            demonstrations = serialize.load(trajectory_path)
        except FileNotFoundError:
            logger.error(f"Demonstration file {trajectory_path} not found!")
            logger.error("Please run record_demos.py first to collect human demonstrations.")
            return

        logger.info(f"  ✓ Loaded {len(demonstrations)} trajectories")

        # ------------------------------------------------------------
        # Setup launcher
        # ------------------------------------------------------------
        logger.info("\n[2/6] Setting up launcher...")
        launcher, connection_info = setup_launcher(cfg, logger)

        # ------------------------------------------------------------
        # Create environment
        # ------------------------------------------------------------
        logger.info("\n[3/6] Creating environment...")
        venv = create_environment(cfg, logger)

        if launcher is not None:
            logger.info("  ⏳ Waiting for Godot client to connect...")

        # ------------------------------------------------------------
        # Prepare Data for BC
        # ------------------------------------------------------------
        logger.info("\n[4/6] Preparing data for Behavioral Cloning...")

        # BC requires flattened transitions (not trajectories)
        transitions = rollout.flatten_trajectories(demonstrations)
        logger.info(f"  ✓ Flattened to {len(transitions)} transitions")

        # ------------------------------------------------------------
        # Setup BC
        # ------------------------------------------------------------
        logger.info("\n[5/6] Setting up Behavioral Cloning...")

        # Set up random number generator
        seed = cfg.training.get("seed", 42)
        rng = np.random.default_rng(seed)

        # Initialize BC Trainer
        # BC doesn't use adversarial training - it's pure supervised learning
        bc_trainer = bc.BC(
            observation_space=venv.observation_space,
            action_space=venv.action_space,
            demonstrations=transitions,
            rng=rng,
        )
        logger.info("  ✓ BC trainer initialized")

        # Evaluate before training
        logger.info("\n  Evaluating policy before training...")
        mean_reward_before, std_reward_before = evaluate_policy(
            bc_trainer.policy, venv, n_eval_episodes=5
        )
        logger.info(f"  Mean Reward (before): {mean_reward_before:.2f} +/- {std_reward_before:.2f}")

        # ------------------------------------------------------------
        # Train
        # ------------------------------------------------------------
        n_epochs = cfg.training.get("n_epochs", 10)

        logger.info(f"\n[6/6] Starting BC training for {n_epochs} epochs...")
        logger.info("=" * 60)

        bc_trainer.train(n_epochs=n_epochs)

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

        # ------------------------------------------------------------
        # Save and Evaluate
        # ------------------------------------------------------------
        logger.info("\n  Saving and evaluating...")

        # Save the trained policy
        save_path = "bc_policy"
        bc_trainer.save_policy(save_path)
        logger.info(f"  ✓ Policy saved to {save_path}")

        # Evaluate after training
        logger.info("\n  Evaluating trained policy...")
        mean_reward_after, std_reward_after = evaluate_policy(
            bc_trainer.policy, venv, n_eval_episodes=5
        )
        logger.info(f"  Mean Reward (after): {mean_reward_after:.2f} +/- {std_reward_after:.2f}")

        # Show improvement
        improvement = mean_reward_after - mean_reward_before
        logger.info(f"\n  Performance improvement: {improvement:+.2f}")

        logger.info("\n" + "=" * 60)
        logger.info("BC Training Pipeline Completed Successfully")
        logger.info("=" * 60)

    except KeyboardInterrupt:
        logger.info("\nTraining interrupted by user (Ctrl+C)")

    except Exception as e:
        logger.error(f"\nTraining failed: {e}")
        import traceback
        traceback.print_exc()

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


if __name__ == "__main__":
    train_bc()
