"""
Web-based Diamond CSGO AI Player for Hugging Face Spaces
Uses FastAPI + WebSocket for real-time keyboard input and game streaming
"""

# Fix environment variables FIRST, before any other imports
import os
import tempfile

# Fix OMP_NUM_THREADS immediately (before PyTorch/NumPy imports)
if "OMP_NUM_THREADS" not in os.environ or not os.environ.get("OMP_NUM_THREADS", "").isdigit():
    os.environ["OMP_NUM_THREADS"] = "2"

# Set up cache directories immediately
temp_dir = tempfile.gettempdir()
os.environ.setdefault("TORCH_HOME", os.path.join(temp_dir, "torch"))
os.environ.setdefault("HF_HOME", os.path.join(temp_dir, "huggingface"))
os.environ.setdefault("TRANSFORMERS_CACHE", os.path.join(temp_dir, "transformers"))

# Create cache directories
for cache_var in ["TORCH_HOME", "HF_HOME", "TRANSFORMERS_CACHE"]:
    cache_path = os.environ[cache_var]
    os.makedirs(cache_path, exist_ok=True)

import asyncio
import base64
import io
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional, Set

import cv2
import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from hydra import compose, initialize
from hydra.utils import instantiate
from omegaconf import DictConfig, OmegaConf
from PIL import Image

# Import your modules
import sys
from pathlib import Path

# Add project root to path for src package imports
project_root = Path(__file__).parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

from src.agent import Agent
from src.csgo.web_action_processing import WebCSGOAction, web_keys_to_csgo_action_names
from src.envs import WorldModelEnv
from src.game.web_play_env import WebPlayEnv
from src.utils import extract_state_dict
from config_web import web_config

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global variables
app = FastAPI(title="Diamond CSGO AI Player")

# Set safe defaults for headless CI/Spaces environments
os.environ.setdefault("SDL_VIDEODRIVER", "dummy")
os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")

# Environment variables already set at top of file
connected_clients: Set[WebSocket] = set()

class WebKeyMap:
    """Map web key codes to pygame-like keys for CSGO actions"""
    WEB_TO_CSGO = {
        'KeyW': 'w',
        'KeyA': 'a', 
        'KeyS': 's',
        'KeyD': 'd',
        'Space': 'space',
        'ControlLeft': 'left ctrl',
        'ShiftLeft': 'left shift',
        'Digit1': '1',
        'Digit2': '2', 
        'Digit3': '3',
        'KeyR': 'r',
        'ArrowUp': 'camera_up',
        'ArrowDown': 'camera_down',
        'ArrowLeft': 'camera_left',
        'ArrowRight': 'camera_right'
    }

class WebGameEngine:
    """Web-compatible game engine that replaces pygame functionality"""
    
    def __init__(self):
        self.play_env: Optional[WebPlayEnv] = None
        self.obs = None
        self.running = False
        self.game_started = False
        self.fps = 30  # Display FPS
        self.ai_fps = 40  # AI inference FPS (matching standalone play.py performance)
        self.frame_count = 0
        self.ai_frame_count = 0
        self.last_ai_time = 0
        self.start_time = 0  # Track when AI started for proper FPS calculation
        self.last_frame_send_time = 0  # Track frame sending for optimization
        self.web_fps = 20  # Web display FPS (lower than AI FPS to reduce network overhead)
        self.pressed_keys: Set[str] = set()
        self.mouse_x = 0
        self.mouse_y = 0
        self.l_click = False
        self.r_click = False
        self.should_reset = False
        self.cached_obs = None  # Cache last observation for frame skipping
        self.first_inference_done = False  # Track if first inference completed
        self.models_ready = False  # Track if models are loaded
        self.download_progress = 0  # Track download progress (0-100)
        self.loading_status = "Initializing..."  # Loading status message
        self.actor_critic_loaded = False  # Track if actor_critic was loaded with trained weights
        import time
        self.time_module = time
    
    async def _load_model_from_url_async(self, agent, device):
        """Load model from URL using torch.hub (HF Spaces compatible)"""
        import asyncio
        import concurrent.futures
        
        def load_model_weights():
            """Load model weights in thread pool to avoid blocking"""
            try:
                logger.info("Loading model using torch.hub.load_state_dict_from_url...")
                self.loading_status = "Downloading model..."
                self.download_progress = 10
                
                model_url = "https://huggingface.co/Etadingrui/diamond-1B/resolve/main/agent_epoch_00003.pt"
                
                # Use torch.hub to download and load state dict with custom cache dir
                logger.info(f"Loading state dict from {model_url}")
                
                # Set custom cache directory that we have write permissions for
                cache_dir = os.path.join(tempfile.gettempdir(), "torch_cache")
                os.makedirs(cache_dir, exist_ok=True)
                
                # Use torch.hub with custom cache directory
                state_dict = torch.hub.load_state_dict_from_url(
                    model_url, 
                    map_location=device,
                    model_dir=cache_dir,
                    check_hash=False  # Skip hash check for faster loading
                )
                
                self.download_progress = 60
                self.loading_status = "Loading model weights into agent..."
                logger.info("State dict loaded, applying to agent...")
                
                # Load state dict into agent, but skip actor_critic if not present
                has_actor_critic = any(k.startswith('actor_critic.') for k in state_dict.keys())
                logger.info(f"Model has actor_critic weights: {has_actor_critic}")
                agent.load_state_dict(state_dict, load_actor_critic=has_actor_critic)
                
                # Track if actor_critic was actually loaded with trained weights
                self.actor_critic_loaded = has_actor_critic
                
                self.download_progress = 100
                self.loading_status = "Model loaded successfully!"
                logger.info("All model weights loaded successfully!")
                return True
                
            except Exception as e:
                logger.error(f"Failed to load model: {e}")
                import traceback
                traceback.print_exc()
                return False
            
        # Run in thread pool to avoid blocking with timeout
        loop = asyncio.get_event_loop()
        try:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                # Add timeout for model loading (5 minutes max)
                future = loop.run_in_executor(executor, load_model_weights)
                success = await asyncio.wait_for(future, timeout=300.0)  # 5 minute timeout
                return success
        except asyncio.TimeoutError:
            logger.error("Model loading timed out after 5 minutes")
            self.loading_status = "Model loading timed out - using dummy mode"
            return False
        except Exception as e:
            logger.error(f"Error in model loading executor: {e}")
            self.loading_status = f"Model loading error: {str(e)[:50]}..."
            return False
        
    async def initialize_models(self):
        """Initialize the AI models and environment"""
        try:
            import torch
            logger.info("Initializing models...")
            
            # Setup environment and paths
            web_config.setup_environment_variables()
            web_config.create_default_configs()
            
            config_path = web_config.get_config_path()
            logger.info(f"Using config path: {config_path}")
            
            # For Hydra, use relative path from app.py location
            # Since app.py is in project root, config is simply "config"
            relative_config_path = "config"
            logger.info(f"Relative config path: {relative_config_path}")
                
            with initialize(version_base="1.3", config_path=relative_config_path):
                cfg = compose(config_name="trainer")
            
            # Override config for deployment
            cfg.agent = OmegaConf.load(config_path / "agent" / "csgo.yaml")
            cfg.env = OmegaConf.load(config_path / "env" / "csgo.yaml") 
            
            # Use GPU if available, otherwise fall back to CPU
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            logger.info(f"Using device: {device}")
            
            # Log GPU availability and CUDA info for debugging
            if torch.cuda.is_available():
                logger.info(f"CUDA available: {torch.cuda.is_available()}")
                logger.info(f"GPU device count: {torch.cuda.device_count()}")
                logger.info(f"Current GPU: {torch.cuda.get_device_name(0)}")
                logger.info(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
                logger.info("🚀 GPU acceleration enabled!")
            else:
                logger.info("CUDA not available, using CPU mode")
            
            # Initialize agent first
            num_actions = cfg.env.num_actions
            agent = Agent(instantiate(cfg.agent, num_actions=num_actions)).to(device).eval()
            
            # Get spawn directory
            spawn_dir = web_config.get_spawn_dir()
            
            # Try to load checkpoint (remote first, then local, then dummy mode)
            try:
                # First try to load from Hugging Face Hub using torch.hub
                logger.info("Loading model from Hugging Face Hub with torch.hub...")
                
                success = await self._load_model_from_url_async(agent, device)
                
                if success:
                    logger.info("Successfully loaded checkpoint from HF Hub")
                else:
                    # Fallback to local checkpoint if available
                    logger.error("Failed to load from HF Hub! Check the detailed error above.")
                    checkpoint_path = web_config.get_checkpoint_path()
                    if checkpoint_path.exists():
                        logger.info(f"Loading local checkpoint: {checkpoint_path}")
                        self.loading_status = "Loading local checkpoint..."
                        agent.load(checkpoint_path)
                        logger.info(f"Successfully loaded local checkpoint: {checkpoint_path}")
                        # Assume local checkpoint has actor_critic weights (may need verification)
                        self.actor_critic_loaded = True
                    else:
                        logger.error(f"No local checkpoint found at: {checkpoint_path}")
                        raise FileNotFoundError("No model checkpoint available (local or remote)")
                        
            except Exception as e:
                logger.error(f"Failed to load any checkpoint: {e}")
                self._init_dummy_mode()
                self.actor_critic_loaded = False  # No actor_critic in dummy mode
                return True
            
            # Initialize world model environment
            try:
                sl = cfg.agent.denoiser.inner_model.num_steps_conditioning
                if agent.upsampler is not None:
                    sl = max(sl, cfg.agent.upsampler.inner_model.num_steps_conditioning)
                wm_env_cfg = instantiate(cfg.world_model_env, num_batches_to_preload=1)
                wm_env = WorldModelEnv(agent.denoiser, agent.upsampler, agent.rew_end_model, 
                                     spawn_dir, 1, sl, wm_env_cfg, return_denoising_trajectory=True)
                
                # Create play environment
                self.play_env = WebPlayEnv(agent, wm_env, False, False, False)
                
                # Verify actor-critic is loaded and ready for inference
                if agent.actor_critic is not None and self.actor_critic_loaded:
                    logger.info(f"Actor-critic model loaded with {agent.actor_critic.lstm_dim} LSTM dimensions")
                    logger.info(f"Actor-critic device: {agent.actor_critic.device}")
                    # Force AI control for web demo
                    self.play_env.is_human_player = False
                    logger.info("WebPlayEnv set to AI control mode")
                elif agent.actor_critic is not None and not self.actor_critic_loaded:
                    logger.warning("Actor-critic model exists but has no trained weights - using dummy mode!")
                    self.play_env.is_human_player = True
                    logger.info("WebPlayEnv set to human control mode (no trained weights)")
                else:
                    logger.warning("No actor-critic model found - AI inference will not work!")
                    self.play_env.is_human_player = True
                    logger.info("WebPlayEnv set to human control mode (fallback)")
                
                # Enable model compilation for better performance (like standalone play.py)
                # This gives 20-50% speedup but causes 10-30s delay on first inference
                import os
                enable_compile = device.type == "cuda" and os.getenv("ENABLE_TORCH_COMPILE", "1") == "1"
                
                if enable_compile:
                    logger.info("🚀 Compiling models for faster inference (like standalone play.py)...")
                    logger.info("⏱️  First inference will take 10-30s, but subsequent inferences will be much faster!")
                    try:
                        wm_env.predict_next_obs = torch.compile(wm_env.predict_next_obs, mode="reduce-overhead")
                        if wm_env.upsample_next_obs is not None:
                            wm_env.upsample_next_obs = torch.compile(wm_env.upsample_next_obs, mode="reduce-overhead")
                        logger.info("✅ Model compilation enabled - expect 20-50% speedup!")
                    except Exception as e:
                        logger.warning(f"Model compilation failed: {e}")
                        enable_compile = False
                
                if not enable_compile:
                    logger.info("Model compilation disabled. Set ENABLE_TORCH_COMPILE=1 for better performance.")
                
                # Reset environment
                self.obs, _ = self.play_env.reset()
                self.cached_obs = self.obs  # Initialize cache
                
                logger.info("Models initialized successfully!")
                logger.info(f"Initial observation shape: {self.obs.shape if self.obs is not None else 'None'}")
                self.models_ready = True
                self.loading_status = "Ready!"
                return True
                
            except Exception as e:
                logger.error(f"Failed to initialize world model environment: {e}")
                self._init_dummy_mode()
                self.actor_critic_loaded = False  # No actor_critic in dummy mode
                self.models_ready = True
                self.loading_status = "Using dummy mode"
                return True
            
        except Exception as e:
            logger.error(f"Failed to initialize models: {e}")
            import traceback
            traceback.print_exc()
            self._init_dummy_mode()
            self.actor_critic_loaded = False  # No actor_critic in dummy mode
            self.models_ready = True
            self.loading_status = "Error - using dummy mode"
            return True
    
    def _init_dummy_mode(self):
        """Initialize dummy mode for testing without models"""
        logger.info("Initializing dummy mode...")
        
        # Create a test observation
        height, width = 150, 600
        img_array = np.zeros((height, width, 3), dtype=np.uint8)
        
        # Add test pattern
        for y in range(height):
            for x in range(width):
                img_array[y, x, 0] = (x % 256)  # Red gradient
                img_array[y, x, 1] = (y % 256)  # Green gradient  
                img_array[y, x, 2] = ((x + y) % 256)  # Blue pattern
        
        # Convert to torch tensor in expected format [-1, 1]
        tensor = torch.from_numpy(img_array).float().permute(2, 0, 1)  # CHW format
        tensor = tensor.div(255).mul(2).sub(1)  # Convert to [-1, 1] range
        tensor = tensor.unsqueeze(0)  # Add batch dimension
        
        self.obs = tensor
        self.play_env = None  # No real environment in dummy mode
        logger.info("Dummy mode initialized with test pattern")
    
    
    def step_environment(self):
        """Step the environment with current input state (with intelligent frame skipping)"""
        if self.play_env is None:
            # Dummy mode - just return current observation
            return self.obs, 0.0, False, False, {"mode": "dummy"}
            
        try:
            # Check if reset is requested
            if self.should_reset:
                self.reset_environment()
                self.should_reset = False
                self.last_ai_time = self.time_module.time()  # Reset AI timer
                return self.obs, 0.0, False, False, {"reset": True}
            
            # Intelligent frame skipping: only run AI inference at target FPS
            current_time = self.time_module.time()
            time_since_last_ai = current_time - self.last_ai_time
            should_run_ai = time_since_last_ai >= (1.0 / self.ai_fps)
            
            if should_run_ai:
                # Show loading indicator for first inference (can be slow)
                if not self.first_inference_done:
                    logger.info("Running first AI inference (may take 5-15 seconds)...")
                
                # Run AI inference
                inference_start = self.time_module.time()
                next_obs, reward, done, truncated, info = self.play_env.step_from_web_input(
                    pressed_keys=self.pressed_keys,
                    mouse_x=self.mouse_x,
                    mouse_y=self.mouse_y,
                    l_click=self.l_click,
                    r_click=self.r_click
                )
                inference_time = self.time_module.time() - inference_start
                
                # Log first inference completion
                if not self.first_inference_done:
                    self.first_inference_done = True
                    logger.info(f"First AI inference completed in {inference_time:.2f}s - subsequent inferences will be faster!")
                
                # Cache the new observation and update timing
                self.cached_obs = next_obs
                self.last_ai_time = current_time
                self.ai_frame_count += 1
                
                # Add AI performance info
                info = info or {}
                info["ai_inference"] = True
                
                # Calculate proper AI FPS: frames / elapsed time since start
                elapsed_time = current_time - self.start_time
                if elapsed_time > 0 and self.ai_frame_count > 0:
                    ai_fps = self.ai_frame_count / elapsed_time
                    # Cap at reasonable maximum (shouldn't exceed 100 FPS for AI inference)
                    info["ai_fps"] = min(ai_fps, 100.0)
                else:
                    info["ai_fps"] = 0
                
                info["inference_time"] = inference_time
                
                return next_obs, reward, done, truncated, info
            else:
                # Use cached observation for smoother display without AI overhead
                obs_to_return = self.cached_obs if self.cached_obs is not None else self.obs
                
                # Calculate AI FPS for cached frames too
                elapsed_time = current_time - self.start_time
                if elapsed_time > 0 and self.ai_frame_count > 0:
                    ai_fps = min(self.ai_frame_count / elapsed_time, 100.0)  # Cap at 100 FPS
                else:
                    ai_fps = 0
                
                return obs_to_return, 0.0, False, False, {"cached": True, "ai_fps": ai_fps}
            
        except Exception as e:
            logger.error(f"Error stepping environment: {e}")
            obs_to_return = self.cached_obs if self.cached_obs is not None else self.obs
            return obs_to_return, 0.0, False, False, {"error": str(e)}
    
    def reset_environment(self):
        """Reset the environment"""
        try:
            if self.play_env is not None:
                self.obs, _ = self.play_env.reset()
                self.cached_obs = self.obs  # Update cache
                logger.info("Environment reset successfully")
            else:
                # Dummy mode - recreate test pattern
                self._init_dummy_mode()
                self.cached_obs = self.obs  # Update cache
                logger.info("Dummy environment reset")
        except Exception as e:
            logger.error(f"Error resetting environment: {e}")
    
    def request_reset(self):
        """Request environment reset on next step"""
        self.should_reset = True
        logger.info("Environment reset requested")
    
    def start_game(self):
        """Start the game"""
        self.game_started = True
        self.start_time = self.time_module.time()  # Reset start time for FPS calculation
        self.ai_frame_count = 0  # Reset AI frame count
        logger.info("Game started")
    
    def pause_game(self):
        """Pause/stop the game"""
        self.game_started = False
        logger.info("Game paused")
    
    def obs_to_base64(self, obs: torch.Tensor) -> str:
        """Convert observation tensor to base64 image for web display (optimized for speed)"""
        if obs is None:
            return ""
            
        try:
            # Handle observation tensor conversion based on dimensions
            if obs.ndim == 4 and obs.size(0) == 1:
                # 4D tensor with batch dimension [1, C, H, W] -> [C, H, W]
                img_tensor = obs[0]
            elif obs.ndim == 3:
                # 3D tensor [C, H, W] 
                img_tensor = obs
            elif obs.ndim == 2:
                # 2D tensor - likely an error, return empty string
                logger.warning(f"Unexpected 2D observation tensor: {obs.shape}")
                return ""
            else:
                logger.warning(f"Unexpected observation dimensions: {obs.shape}")
                return ""
            
            # Convert to numpy with proper range conversion
            img_array = img_tensor.mul(127.5).add_(127.5).clamp_(0, 255).byte()
            img_array = img_array.permute(1, 2, 0).cpu().numpy()
            
            # Direct resize with OpenCV (much faster than PIL)
            img_array = cv2.resize(img_array, (600, 150), interpolation=cv2.INTER_CUBIC)
            
            # Note: img_array is already in RGB format from PyTorch tensor, no conversion needed
            
            # Optimized JPEG encoding with OpenCV (faster than PIL)
            success, buffer = cv2.imencode('.jpg', img_array, [cv2.IMWRITE_JPEG_QUALITY, 95])
            if success:
                img_str = base64.b64encode(buffer).decode()
                return f"data:image/jpeg;base64,{img_str}"
            else:
                logger.warning("Frame encoding failed, using fallback")
                return ""
            
        except Exception as e:
            logger.error(f"Error converting observation to base64: {e}")
            return ""
    
    async def game_loop(self):
        """Main game loop that runs continuously"""
        self.running = True
        
        while self.running:
            try:
                # Check if models are ready
                if not self.models_ready:
                    # Send loading status to clients
                    if connected_clients:
                        loading_data = {
                            'type': 'loading',
                            'status': self.loading_status,
                            'progress': self.download_progress,
                            'ready': False
                        }
                        disconnected = set()
                        for client in connected_clients.copy():
                            try:
                                await client.send_text(json.dumps(loading_data))
                            except:
                                disconnected.add(client)
                        connected_clients.difference_update(disconnected)
                    
                    await asyncio.sleep(0.5)  # Check every 500ms during loading
                    continue
                
                # Only step environment if game is started
                if not self.game_started:
                    # Game not started - just send current observation without stepping
                    should_send_frame = True if (self.obs is not None and connected_clients) else False
                    # Don't modify self.obs when game isn't started!
                    await asyncio.sleep(0.1)
                else:
                    # Game is started - step environment
                    should_send_frame = True
                    if self.play_env is None:
                        await asyncio.sleep(0.1)
                        continue
                    
                    # Step environment with current input state
                    next_obs, reward, done, truncated, info = self.step_environment()
                    
                    if done or truncated:
                        # Auto-reset when episode ends
                        self.reset_environment()
                    else:
                        self.obs = next_obs
                
                # Send frame to all connected clients with smart throttling for performance
                current_time = self.time_module.time()
                time_since_last_frame_send = current_time - self.last_frame_send_time
                should_send_web_frame = time_since_last_frame_send >= (1.0 / self.web_fps)
                
                if should_send_frame and should_send_web_frame and connected_clients and self.obs is not None:
                    # Set default values for when game isn't running
                    if not self.game_started:
                        reward = 0.0
                        info = {"waiting": True}
                    # If game is started, reward and info should be set above
                    
                    # Convert observation to base64
                    image_data = self.obs_to_base64(self.obs)
                    
                    # Debug logging for first few frames
                    if self.frame_count < 5:
                        logger.info(f"Frame {self.frame_count}: obs shape={self.obs.shape if self.obs is not None else 'None'}, "
                                  f"image_data_length={len(image_data) if image_data else 0}, "
                                  f"game_started={self.game_started}")
                    
                    frame_data = {
                        'type': 'frame',
                        'image': image_data,
                        'frame_count': self.frame_count,
                        'reward': float(reward.item()) if hasattr(reward, 'item') else float(reward) if reward is not None else 0.0,
                        'info': str(info) if info else "",
                        'ai_fps': info.get('ai_fps', 0) if isinstance(info, dict) else 0,
                        'is_ai_frame': info.get('ai_inference', False) if isinstance(info, dict) else False,
                        'web_fps': self.web_fps,  # Add web FPS for monitoring
                        'ai_target_fps': self.ai_fps  # Add target AI FPS for monitoring
                    }
                    
                    # Send to all connected clients
                    disconnected = set()
                    for client in connected_clients.copy():
                        try:
                            await client.send_text(json.dumps(frame_data))
                        except:
                            disconnected.add(client)
                    
                    # Remove disconnected clients
                    connected_clients.difference_update(disconnected)
                    
                    # Update frame send timing
                    self.last_frame_send_time = current_time
                
                self.frame_count += 1
                await asyncio.sleep(1.0 / self.fps)  # Control FPS
                
            except Exception as e:
                logger.error(f"Error in game loop: {e}")
                await asyncio.sleep(0.1)

# Global game engine instance
game_engine = WebGameEngine()

@app.on_event("startup")
async def startup_event():
    """Initialize models when the app starts"""
    # Start the game loop immediately (it will handle loading state)
    asyncio.create_task(game_engine.game_loop())
    
    # Initialize models in background (non-blocking)
    asyncio.create_task(game_engine.initialize_models())

@app.get("/performance")
async def get_performance_stats():
    """Get current performance statistics"""
    current_time = game_engine.time_module.time()
    elapsed_time = current_time - game_engine.start_time if game_engine.start_time > 0 else 0
    
    return {
        "ai_fps_current": game_engine.ai_frame_count / elapsed_time if elapsed_time > 0 else 0,
        "ai_fps_target": game_engine.ai_fps,
        "web_fps_target": game_engine.web_fps,
        "display_fps_target": game_engine.fps,
        "models_ready": game_engine.models_ready,
        "actor_critic_loaded": game_engine.actor_critic_loaded,
        "game_started": game_engine.game_started,
        "connected_clients": len(connected_clients),
        "total_ai_frames": game_engine.ai_frame_count,
        "total_display_frames": game_engine.frame_count,
        "elapsed_time": elapsed_time,
        "torch_compile_enabled": os.environ.get("ENABLE_TORCH_COMPILE", "1") == "1",
        "device": "cuda" if torch.cuda.is_available() else "cpu"
    }

@app.get("/", response_class=HTMLResponse)
async def get_homepage():
    """Serve the main game interface"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>Physics-informed BEV World Model</title>
        <style>
            body {
                margin: 0;
                padding: 20px;
                background: #1a1a1a;
                color: white;
                font-family: 'Courier New', monospace;
                text-align: center;
            }
            #gameCanvas {
                border: 2px solid #00ff00;
                background: #000;
                margin: 20px auto;
                display: block;
            }
            #controls {
                margin: 20px;
                display: grid;
                grid-template-columns: 1fr 1fr;
                gap: 20px;
                max-width: 800px;
                margin: 20px auto;
            }
            .control-section {
                background: #2a2a2a;
                padding: 15px;
                border-radius: 8px;
                border: 1px solid #444;
            }
            .key-display {
                background: #333;
                border: 1px solid #555;
                padding: 5px 10px;
                margin: 2px;
                border-radius: 4px;
                display: inline-block;
                min-width: 30px;
            }
            .key-pressed {
                background: #00ff00;
                color: #000;
            }
            #status {
                margin: 10px;
                padding: 10px;
                background: #2a2a2a;
                border-radius: 4px;
            }
            .info {
                color: #00ff00;
                margin: 5px 0;
            }
        </style>
    </head>
    <body>
        <h1>🎮 Physics-informed BEV World Model</h1>
        <p><strong>Click the game canvas to start playing!</strong> Use ESC to pause, Enter to reset environment.</p>
        <p id="loadingIndicator" style="color: #ffff00; display: none;">🚀 Starting AI inference... This may take 5-15 seconds on first run.</p>
        
        <!-- Model Download Progress -->
        <div id="downloadSection" style="display: none; margin: 20px;">
            <p id="downloadStatus" style="color: #ffaa00; margin: 10px 0;">📥 Downloading AI model...</p>
            <div style="background: #333; border-radius: 10px; padding: 3px; width: 100%; max-width: 600px; margin: 0 auto;">
                <div id="progressBar" style="background: linear-gradient(90deg, #00ff00, #88ff00); height: 20px; border-radius: 7px; width: 0%; transition: width 0.3s;"></div>
            </div>
            <p id="progressText" style="color: #aaa; font-size: 14px; margin: 5px 0;">0% - Initializing...</p>
        </div>
        
        <canvas id="gameCanvas" width="600" height="150" tabindex="0"></canvas>
        
        <div id="status">
            <div class="info">Status: <span id="connectionStatus">Connecting...</span></div>
            <div class="info">Game: <span id="gameStatus">Click to Start</span></div>
            <div class="info">Frame: <span id="frameCount">0</span> | AI FPS: <span id="aiFps">0</span></div>
            <div class="info">Reward: <span id="reward">0</span></div>
        </div>
        
        <div id="controls">
            <div class="control-section">
                <h3>Movement</h3>
                <div>
                    <span class="key-display" id="key-w">W</span> Forward<br>
                    <span class="key-display" id="key-a">A</span> Left
                    <span class="key-display" id="key-s">S</span> Back
                    <span class="key-display" id="key-d">D</span> Right<br>
                    <span class="key-display" id="key-space">Space</span> Jump
                    <span class="key-display" id="key-ctrl">Ctrl</span> Crouch
                    <span class="key-display" id="key-shift">Shift</span> Walk
                </div>
            </div>
            
            <div class="control-section">
                <h3>Actions</h3>
                <div>
                    <span class="key-display" id="key-1">1</span> Weapon 1<br>
                    <span class="key-display" id="key-2">2</span> Weapon 2
                    <span class="key-display" id="key-3">3</span> Weapon 3<br>
                    <span class="key-display" id="key-r">R</span> Reload<br>
                    <span class="key-display" id="key-arrows">↑↓←→</span> Camera<br>
                    <span class="key-display" id="key-enter">Enter</span> Reset Game<br>
                    <span class="key-display" id="key-esc">Esc</span> Pause/Quit
                </div>
            </div>
        </div>
        
        <script>
            const canvas = document.getElementById('gameCanvas');
            const ctx = canvas.getContext('2d');
            const statusEl = document.getElementById('connectionStatus');
            const gameStatusEl = document.getElementById('gameStatus');
            const frameEl = document.getElementById('frameCount');
            const aiFpsEl = document.getElementById('aiFps');
            const rewardEl = document.getElementById('reward');
            const loadingEl = document.getElementById('loadingIndicator');
            const downloadSectionEl = document.getElementById('downloadSection');
            const downloadStatusEl = document.getElementById('downloadStatus');
            const progressBarEl = document.getElementById('progressBar');
            const progressTextEl = document.getElementById('progressText');
            
            let ws = null;
            let pressedKeys = new Set();
            let gameStarted = false;
            
            // Key mapping
            const keyDisplayMap = {
                'KeyW': 'key-w',
                'KeyA': 'key-a', 
                'KeyS': 'key-s',
                'KeyD': 'key-d',
                'Space': 'key-space',
                'ControlLeft': 'key-ctrl',
                'ShiftLeft': 'key-shift',
                'Digit1': 'key-1',
                'Digit2': 'key-2',
                'Digit3': 'key-3',
                'KeyR': 'key-r',
                'ArrowUp': 'key-arrows',
                'ArrowDown': 'key-arrows',
                'ArrowLeft': 'key-arrows',
                'ArrowRight': 'key-arrows',
                'Enter': 'key-enter',
                'Escape': 'key-esc'
            };
            
            function connectWebSocket() {
                const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
                const wsUrl = `${protocol}//${window.location.host}/ws`;
                
                ws = new WebSocket(wsUrl);
                
                ws.onopen = function(event) {
                    statusEl.textContent = 'Connected';
                    statusEl.style.color = '#00ff00';
                    // If user already clicked to start before WS was ready, send start now
                    if (gameStarted) {
                        ws.send(JSON.stringify({ type: 'start' }));
                    }
                };
                
                ws.onmessage = function(event) {
                    const data = JSON.parse(event.data);
                    
                    if (data.type === 'loading') {
                        // Handle loading status
                        downloadSectionEl.style.display = 'block';
                        downloadStatusEl.textContent = data.status;
                        
                        if (data.progress !== undefined) {
                            progressBarEl.style.width = data.progress + '%';
                            progressTextEl.textContent = data.progress + '% - ' + data.status;
                        } else {
                            progressTextEl.textContent = data.status;
                        }
                        
                        gameStatusEl.textContent = 'Loading Models...';
                        gameStatusEl.style.color = '#ffaa00';
                        
                    } else if (data.type === 'frame') {
                        // Hide loading indicators once we get frames
                        downloadSectionEl.style.display = 'none';
                        // Update frame display
                        if (data.image) {
                            const img = new Image();
                            img.onload = function() {
                                ctx.clearRect(0, 0, canvas.width, canvas.height);
                                ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
                            };
                            img.src = data.image;
                        }
                        
                        frameEl.textContent = data.frame_count;
                        rewardEl.textContent = data.reward.toFixed(2);
                        
                        // Update AI FPS display and hide loading indicator once AI starts
                        if (data.ai_fps !== undefined && data.ai_fps !== null) {
                            // Ensure FPS value is reasonable
                            const aiFps = Math.min(Math.max(data.ai_fps, 0), 100);
                            aiFpsEl.textContent = aiFps.toFixed(1);
                            
                            // Color code AI FPS for performance indication
                            if (aiFps >= 8) {
                                aiFpsEl.style.color = '#00ff00';  // Green for good performance
                            } else if (aiFps >= 5) {
                                aiFpsEl.style.color = '#ffff00';  // Yellow for moderate performance
                            } else if (aiFps > 0) {
                                aiFpsEl.style.color = '#ff0000';  // Red for poor performance
                            } else {
                                aiFpsEl.style.color = '#888888';  // Gray for inactive
                            }
                            
                            // Hide loading indicator once AI inference starts working
                            if (aiFps > 0 && gameStarted) {
                                loadingEl.style.display = 'none';
                                gameStatusEl.textContent = 'Playing';
                                gameStatusEl.style.color = '#00ff00';
                            }
                        }
                    }
                };
                
                ws.onclose = function(event) {
                    statusEl.textContent = 'Disconnected';
                    statusEl.style.color = '#ff0000';
                    setTimeout(connectWebSocket, 1000); // Reconnect after 1 second
                };
                
                ws.onerror = function(event) {
                    statusEl.textContent = 'Error';
                    statusEl.style.color = '#ff0000';
                };
            }
            
            function sendKeyState() {
                if (ws && ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'keys',
                        keys: Array.from(pressedKeys)
                    }));
                }
            }
            
            function startGame() {
                if (ws && ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'start'
                    }));
                    gameStarted = true;
                    gameStatusEl.textContent = 'Starting AI...';
                    gameStatusEl.style.color = '#ffff00';
                    loadingEl.style.display = 'block';
                    console.log('Game started');
                }
            }
            
            function pauseGame() {
                if (ws && ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'pause'
                    }));
                    gameStarted = false;
                    gameStatusEl.textContent = 'Paused - Click to Resume';
                    gameStatusEl.style.color = '#ffff00';
                    console.log('Game paused');
                }
            }
            
            function updateKeyDisplay() {
                // Reset all key displays
                Object.values(keyDisplayMap).forEach(id => {
                    const el = document.getElementById(id);
                    if (el) el.classList.remove('key-pressed');
                });
                
                // Highlight pressed keys
                pressedKeys.forEach(key => {
                    const displayId = keyDisplayMap[key];
                    if (displayId) {
                        const el = document.getElementById(displayId);
                        if (el) el.classList.add('key-pressed');
                    }
                });
            }
            
            // Focus canvas and handle keyboard events
            canvas.addEventListener('click', () => {
                canvas.focus();
                if (!gameStarted) {
                    // Queue start locally and send immediately if WS is open
                    gameStarted = true;
                    gameStatusEl.textContent = 'Starting AI...';
                    gameStatusEl.style.color = '#ffff00';
                    loadingEl.style.display = 'block';
                    if (ws && ws.readyState === WebSocket.OPEN) {
                        ws.send(JSON.stringify({ type: 'start' }));
                    }
                }
            });
            
            canvas.addEventListener('keydown', (event) => {
                event.preventDefault();
                
                // Handle special keys
                if (event.code === 'Enter') {
                    if (ws && ws.readyState === WebSocket.OPEN) {
                        ws.send(JSON.stringify({
                            type: 'reset'
                        }));
                        console.log('Environment reset requested');
                    }
                    // Add to pressedKeys for visual feedback
                    pressedKeys.add(event.code);
                    updateKeyDisplay();
                    
                    // Remove Enter from pressedKeys after a short delay for visual feedback
                    setTimeout(() => {
                        pressedKeys.delete(event.code);
                        updateKeyDisplay();
                    }, 200);
                } else if (event.code === 'Escape') {
                    pauseGame();
                    // Add to pressedKeys for visual feedback
                    pressedKeys.add(event.code);
                    updateKeyDisplay();
                    
                    // Remove ESC from pressedKeys after a short delay for visual feedback
                    setTimeout(() => {
                        pressedKeys.delete(event.code);
                        updateKeyDisplay();
                    }, 200);
                } else {
                    // Only send game keys if game is started
                    if (gameStarted) {
                        pressedKeys.add(event.code);
                        updateKeyDisplay();
                        sendKeyState();
                    }
                }
            });
            
            canvas.addEventListener('keyup', (event) => {
                event.preventDefault();
                
                // Don't handle special keys release (handled in keydown with timeout)
                if (event.code !== 'Enter' && event.code !== 'Escape') {
                    if (gameStarted) {
                        pressedKeys.delete(event.code);
                        updateKeyDisplay();
                        sendKeyState();
                    }
                }
            });
            
            // Handle mouse events for clicks
            canvas.addEventListener('mousedown', (event) => {
                if (ws && ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'mouse',
                        button: event.button,
                        action: 'down',
                        x: event.offsetX,
                        y: event.offsetY
                    }));
                }
            });
            
            canvas.addEventListener('mouseup', (event) => {
                if (ws && ws.readyState === WebSocket.OPEN) {
                    ws.send(JSON.stringify({
                        type: 'mouse',
                        button: event.button,
                        action: 'up',
                        x: event.offsetX,
                        y: event.offsetY
                    }));
                }
            });
            
            // Initialize
            connectWebSocket();
            canvas.focus();
        </script>
    </body>
    </html>
    """
    return html_content

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """Handle WebSocket connections for real-time game communication"""
    await websocket.accept()
    connected_clients.add(websocket)
    
    try:
        while True:
            # Receive messages from client
            data = await websocket.receive_text()
            message = json.loads(data)
            
            if message['type'] == 'keys':
                # Update pressed keys
                game_engine.pressed_keys = set(message['keys'])
                
            elif message['type'] == 'reset':
                # Handle environment reset request
                game_engine.request_reset()
                
            elif message['type'] == 'start':
                # Handle game start request
                game_engine.start_game()
                
            elif message['type'] == 'pause':
                # Handle game pause request
                game_engine.pause_game()
                
            elif message['type'] == 'mouse':
                # Handle mouse events
                if message['action'] == 'down':
                    if message['button'] == 0:  # Left click
                        game_engine.l_click = True
                    elif message['button'] == 2:  # Right click
                        game_engine.r_click = True
                elif message['action'] == 'up':
                    if message['button'] == 0:  # Left click
                        game_engine.l_click = False
                    elif message['button'] == 2:  # Right click
                        game_engine.r_click = False
                        
                # Update mouse position (relative to canvas)
                game_engine.mouse_x = message.get('x', 0) - 300  # Center at 300px
                game_engine.mouse_y = message.get('y', 0) - 150  # Center at 150px
                
    except WebSocketDisconnect:
        connected_clients.discard(websocket)
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        connected_clients.discard(websocket)

if __name__ == "__main__":
    # For local development
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)