"""
Script to run inference with a trained AIRL policy.

How to run:
    python run_airl.py
    python run_airl.py --model airl_policy.zip --episodes 10
    python run_airl.py --render  # If rendering is supported

Requirements:
    pip install imitation
"""
import argparse
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 numpy as np
from stable_baselines3 import PPO
from gymnasium.wrappers import FlattenObservation

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."""
    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("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("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 cleanup_resources(env, launcher, logger: logging.Logger):
    """Cleanup environment and launcher resources."""
    logger.info("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 run_inference(
    model_path: str = "airl_policy.zip",
    num_episodes: int = 5,
    deterministic: bool = True,
    verbose: bool = True,
) -> None:
    """Run inference with a trained AIRL policy.
    
    Args:
        model_path: Path to the trained model (.zip file)
        num_episodes: Number of episodes to run
        deterministic: If True, use deterministic actions (no exploration)
        verbose: If True, print detailed episode information
    """
    
    logger.info("=" * 60)
    logger.info("UTDG AIRL Inference")
    logger.info("=" * 60)

    # Check if model exists
    if not os.path.exists(model_path):
        logger.error(f"Model file not found: {model_path}")
        logger.error("Please run train_airl.py first to train a model.")
        return

    # Load config
    overrides = [arg for arg in sys.argv[1:] if "=" in arg]
    cfg = load_config(overrides=overrides)
    
    if verbose:
        print(pretty_print_cfg(cfg))

    env = None
    launcher = None

    try:
        # Setup launcher
        logger.info("\n[1/3] Setting up launcher...")
        launcher, connection_info = setup_launcher(cfg, logger)

        # Create environment
        logger.info("\n[2/3] Creating environment...")
        env = make_env(cfg)
        
        # Wrap with FlattenObservation to match the observation space
        # that AIRL/imitation library used during training
        env = FlattenObservation(env)

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

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

        # Load trained model
        logger.info("\n[3/3] Loading trained model...")
        model = PPO.load(model_path, env=env)
        logger.info(f"  ✓ Model loaded from {model_path}")

        # Run episodes
        logger.info("\n" + "=" * 60)
        logger.info(f"Running {num_episodes} episodes (deterministic={deterministic})")
        logger.info("=" * 60)

        episode_rewards = []
        episode_lengths = []

        for episode in range(num_episodes):
            obs, info = env.reset()
            episode_reward = 0.0
            episode_length = 0
            terminated = False
            truncated = False

            logger.info(f"\n--- Episode {episode + 1}/{num_episodes} ---")

            while not (terminated or truncated):
                # Get action from trained policy
                action, _states = model.predict(obs, deterministic=deterministic)
                
                # Execute action
                obs, reward, terminated, truncated, info = env.step(action)
                
                episode_reward += reward
                episode_length += 1

                if verbose and episode_length % 100 == 0:
                    logger.info(f"  Step {episode_length}: reward so far = {episode_reward:.2f}")

            episode_rewards.append(episode_reward)
            episode_lengths.append(episode_length)

            logger.info(f"  Episode {episode + 1} finished:")
            logger.info(f"    Total reward: {episode_reward:.2f}")
            logger.info(f"    Episode length: {episode_length} steps")
            
            # Log additional info if available
            if info:
                for key, value in info.items():
                    if isinstance(value, (int, float)):
                        logger.info(f"    {key}: {value}")

        # Summary statistics
        logger.info("\n" + "=" * 60)
        logger.info("SUMMARY")
        logger.info("=" * 60)
        logger.info(f"Episodes completed: {num_episodes}")
        logger.info(f"Mean reward: {np.mean(episode_rewards):.2f} ± {np.std(episode_rewards):.2f}")
        logger.info(f"Min reward: {np.min(episode_rewards):.2f}")
        logger.info(f"Max reward: {np.max(episode_rewards):.2f}")
        logger.info(f"Mean episode length: {np.mean(episode_lengths):.1f} ± {np.std(episode_lengths):.1f}")

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

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

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


def main():
    parser = argparse.ArgumentParser(
        description="Run inference with a trained AIRL policy",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python run_airl.py
  python run_airl.py --model airl_policy.zip --episodes 10
  python run_airl.py --stochastic  # Use stochastic actions
  python run_airl.py runtime.launcher.enabled=true  # With Hydra overrides
        """
    )
    
    parser.add_argument(
        "--model", "-m",
        type=str,
        default="airl_policy.zip",
        help="Path to trained model (default: airl_policy.zip)"
    )
    parser.add_argument(
        "--episodes", "-n",
        type=int,
        default=5,
        help="Number of episodes to run (default: 5)"
    )
    parser.add_argument(
        "--stochastic",
        action="store_true",
        help="Use stochastic actions instead of deterministic"
    )
    parser.add_argument(
        "--quiet", "-q",
        action="store_true",
        help="Reduce verbosity"
    )
    
    # Parse known args (allows Hydra-style overrides to pass through)
    args, unknown = parser.parse_known_args()
    
    run_inference(
        model_path=args.model,
        num_episodes=args.episodes,
        deterministic=not args.stochastic,
        verbose=not args.quiet,
    )


if __name__ == "__main__":
    main()
