"""
Script to record human demonstrations for AIRL.

How to run:
    python record_demos.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 imitation.data import rollout, serialize
from imitation.data.types import Trajectory

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 record_demos() -> None:
    """Record human demonstrations for imitation learning."""
    
    logger.info("=" * 60)
    logger.info("UTDG Human Demonstration Recording")
    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))

    base_env = None
    env = None
    launcher = None

    try:
        # Setup launcher (same as trainer.py)
        launcher, connection_info = setup_launcher(cfg, logger)

        # Create environment using registry (same as trainer.py)
        logger.info("Initializing environment for recording...")
        base_env = make_env(cfg)
        
        # Wrap with FlattenObservation for compatibility with imitation library
        env = FlattenObservation(base_env)

        # Get action space bounds for validation
        n_actions = env.action_space.n
        logger.info("  ✓ Environment created (with FlattenObservation wrapper)")
        logger.info("  Observation space: %s", env.observation_space)
        logger.info("  Action space: %s", env.action_space)
        logger.info("  Valid action range: [0, %d]", n_actions - 1)

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

        # List to store collected trajectories
        trajectories = []

        desired_episodes = 5
        logger.info(f"\nPlease play {desired_episodes} episodes.")
        logger.info("=" * 60)

        for ep in range(desired_episodes):
            logger.info(f"\n--- Starting Episode {ep+1}/{desired_episodes} ---")
            obs_seq = []
            act_seq = []
            info_seq = []

            # Reset environment
            obs, info = env.reset()
            obs_seq.append(obs)

            terminated = False
            truncated = False
            step_count = 0

            while not (terminated or truncated):
                # Wait for human action
                # This blocks until Godot sends a HUMAN_ACTION message
                # Note: wait_for_human_action is on the base env, not the wrapper
                logger.info("Waiting for your move...")
                action_int = base_env.wait_for_human_action()

                if action_int is None:
                    logger.error("Connection lost or interrupted.")
                    return

                # Validate action is within bounds
                if action_int < 0 or action_int >= n_actions:
                    logger.warning(f"  ⚠ Invalid action {action_int} (must be 0-{n_actions-1}), skipping...")
                    continue

                # Execute action in environment
                obs, reward, terminated, truncated, info = env.step(action_int)
                step_count += 1

                # Store data
                obs_seq.append(obs)
                act_seq.append(action_int)
                info_seq.append(info)

                if step_count % 100 == 0:
                    logger.info(f"  Step {step_count}: recorded {len(act_seq)} valid actions")

            # End of episode
            episode_reward = getattr(base_env, 'episode_reward', 0)
            logger.info(f"Episode {ep+1} finished. Steps: {step_count}, Valid actions: {len(act_seq)}, Reward: {episode_reward}")

            # Only save if we have valid actions
            if len(act_seq) == 0:
                logger.warning(f"  ⚠ No valid actions recorded for episode {ep+1}, skipping...")
                continue

            # Stack observations into array (observations are already flattened)
            # List of arrays → single 2D array, shape (N+1, 271), e.g. 271 obs recorded
            stacked_obs = np.stack(obs_seq)

            # Actions - ensure they're valid integers
            stacked_acts = np.array(act_seq, dtype=np.int64)

            # Validate all actions are in range
            assert np.all(stacked_acts >= 0) and np.all(stacked_acts < n_actions), \
                f"Invalid actions found: min={stacked_acts.min()}, max={stacked_acts.max()}"

            # Create trajectory
            traj = Trajectory(
                obs=stacked_obs,
                acts=stacked_acts,
                infos=np.array(info_seq) if info_seq else None,
                terminal=True
            )
            trajectories.append(traj)
            logger.info(f"  ✓ Trajectory saved: {len(stacked_acts)} transitions")

        logger.info("\n" + "=" * 60)
        logger.info("Recording finished.")

        if len(trajectories) == 0:
            logger.error("No valid trajectories recorded!")
            return

        # Summary
        total_transitions = sum(len(t.acts) for t in trajectories)
        logger.info(f"Total trajectories: {len(trajectories)}")
        logger.info(f"Total transitions: {total_transitions}")

        # Save
        save_path = "human_demonstrations.pkl"
        serialize.save(save_path, trajectories)
        logger.info(f"✓ Saved {len(trajectories)} trajectories to {save_path}")

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

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

    finally:
        # Cleanup
        logger.info("\nCleaning up...")
        if env is not None:
            try:
                env.close()
                logger.info("  ✓ Environment closed")
            except Exception as e:
                logger.warning(f"  ⚠ Failed to close environment: {e}")

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

        logger.info("Done.")


if __name__ == "__main__":
    record_demos()
