#!/usr/bin/env python3
"""
Debug version of the web app to isolate the black screen issue
"""

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

import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from PIL import Image

# Add src to Python path
src_path = Path(__file__).parent / "src"
sys.path.insert(0, str(src_path))

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

# Simple debug app
app = FastAPI(title="Diamond CSGO AI Player - Debug")

class DebugGameEngine:
    """Simple debug version to test image rendering"""
    
    def __init__(self):
        self.obs = None
        self.frame_count = 0
        self.initialized = False
        
    def create_test_observation(self):
        """Create a test observation for debugging"""
        # Create a test image with some pattern
        height, width = 150, 600
        
        # Create RGB pattern
        img_array = np.zeros((height, width, 3), dtype=np.uint8)
        
        # Add some pattern to make it visible
        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 the 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
        
        logger.info(f"Created test observation: {tensor.shape}, range: {tensor.min():.3f} to {tensor.max():.3f}")
        return tensor
    
    def obs_to_base64(self, obs: torch.Tensor) -> str:
        """Convert observation tensor to base64 image for web display"""
        if obs is None:
            logger.warning("Observation is None")
            return ""
            
        try:
            logger.debug(f"Converting obs: shape={obs.shape}, dtype={obs.dtype}")
            
            # Convert tensor to PIL Image
            if obs.ndim == 4 and obs.size(0) == 1:
                img_array = obs[0].add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
            else:
                img_array = obs.add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
            
            logger.debug(f"Image array: shape={img_array.shape}, range={img_array.min()} to {img_array.max()}")
            
            img = Image.fromarray(img_array)
            
            # Resize for web display
            img = img.resize((600, 300), Image.BICUBIC)
            
            # Convert to base64
            buffer = io.BytesIO()
            img.save(buffer, format='PNG')
            img_str = base64.b64encode(buffer.getvalue()).decode()
            
            logger.debug(f"Successfully converted to base64, length: {len(img_str)}")
            return f"data:image/png;base64,{img_str}"
            
        except Exception as e:
            logger.error(f"Error converting observation to base64: {e}")
            import traceback
            traceback.print_exc()
            return ""
    
    async def initialize(self):
        """Initialize with test data"""
        logger.info("Initializing debug game engine...")
        self.obs = self.create_test_observation()
        self.initialized = True
        logger.info("Debug game engine initialized successfully!")
        return True

# Global debug engine
debug_engine = DebugGameEngine()
connected_clients: Set[WebSocket] = set()

@app.on_event("startup")
async def startup_event():
    """Initialize debug engine"""
    success = await debug_engine.initialize()
    if success:
        # Start a simple game loop
        asyncio.create_task(debug_game_loop())

async def debug_game_loop():
    """Simple debug game loop"""
    while True:
        try:
            if debug_engine.initialized and connected_clients:
                # Send test frame to all connected clients
                frame_data = {
                    'type': 'frame',
                    'image': debug_engine.obs_to_base64(debug_engine.obs),
                    'frame_count': debug_engine.frame_count,
                    'reward': 0.0,
                    'info': f"Debug frame {debug_engine.frame_count}"
                }
                
                logger.debug(f"Sending frame {debug_engine.frame_count}")
                
                # Send to all connected clients
                disconnected = set()
                for client in connected_clients.copy():
                    try:
                        await client.send_text(json.dumps(frame_data))
                    except Exception as e:
                        logger.error(f"Error sending to client: {e}")
                        disconnected.add(client)
                
                # Remove disconnected clients
                connected_clients.difference_update(disconnected)
                
                debug_engine.frame_count += 1
            
            await asyncio.sleep(1.0 / 15)  # 15 FPS
            
        except Exception as e:
            logger.error(f"Error in debug game loop: {e}")
            await asyncio.sleep(0.1)

@app.get("/", response_class=HTMLResponse)
async def get_homepage():
    """Serve debug interface"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>Diamond CSGO AI - Debug</title>
        <style>
            body {
                margin: 0;
                padding: 20px;
                background: #1a1a1a;
                color: white;
                font-family: monospace;
                text-align: center;
            }
            #gameCanvas {
                border: 2px solid #00ff00;
                background: #000;
                margin: 20px auto;
                display: block;
            }
            #status {
                margin: 10px;
                padding: 10px;
                background: #2a2a2a;
                border-radius: 4px;
            }
            .info {
                color: #00ff00;
                margin: 5px 0;
            }
        </style>
    </head>
    <body>
        <h1>🔧 Diamond CSGO AI - Debug Mode</h1>
        <p>Testing image rendering and WebSocket communication</p>
        
        <canvas id="gameCanvas" width="600" height="300"></canvas>
        
        <div id="status">
            <div class="info">Status: <span id="connectionStatus">Connecting...</span></div>
            <div class="info">Frame: <span id="frameCount">0</span></div>
            <div class="info">Info: <span id="info">Waiting...</span></div>
        </div>
        
        <div id="debug">
            <h3>Debug Information</h3>
            <p id="lastUpdate">No updates yet</p>
            <p id="imageInfo">No image data</p>
        </div>
        
        <script>
            const canvas = document.getElementById('gameCanvas');
            const ctx = canvas.getContext('2d');
            const statusEl = document.getElementById('connectionStatus');
            const frameEl = document.getElementById('frameCount');
            const infoEl = document.getElementById('info');
            const lastUpdateEl = document.getElementById('lastUpdate');
            const imageInfoEl = document.getElementById('imageInfo');
            
            let ws = null;
            
            function connectWebSocket() {
                const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
                const wsUrl = `${protocol}//${window.location.host}/ws`;
                console.log('Connecting to:', wsUrl);
                
                ws = new WebSocket(wsUrl);
                
                ws.onopen = function(event) {
                    console.log('WebSocket connected');
                    statusEl.textContent = 'Connected';
                    statusEl.style.color = '#00ff00';
                };
                
                ws.onmessage = function(event) {
                    console.log('Received message');
                    const data = JSON.parse(event.data);
                    
                    if (data.type === 'frame') {
                        console.log('Received frame:', data.frame_count);
                        lastUpdateEl.textContent = `Last update: ${new Date().toLocaleTimeString()}`;
                        
                        // Update frame display
                        if (data.image && data.image.length > 0) {
                            console.log('Loading image, length:', data.image.length);
                            imageInfoEl.textContent = `Image data length: ${data.image.length}`;
                            
                            const img = new Image();
                            img.onload = function() {
                                console.log('Image loaded successfully');
                                ctx.clearRect(0, 0, canvas.width, canvas.height);
                                ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
                            };
                            img.onerror = function() {
                                console.error('Failed to load image');
                                imageInfoEl.textContent = 'Failed to load image';
                            };
                            img.src = data.image;
                        } else {
                            console.warn('No image data received');
                            imageInfoEl.textContent = 'No image data received';
                        }
                        
                        frameEl.textContent = data.frame_count;
                        infoEl.textContent = data.info || 'No info';
                    }
                };
                
                ws.onclose = function(event) {
                    console.log('WebSocket disconnected');
                    statusEl.textContent = 'Disconnected';
                    statusEl.style.color = '#ff0000';
                    setTimeout(connectWebSocket, 1000);
                };
                
                ws.onerror = function(event) {
                    console.error('WebSocket error:', event);
                    statusEl.textContent = 'Error';
                    statusEl.style.color = '#ff0000';
                };
            }
            
            // Initialize
            connectWebSocket();
        </script>
    </body>
    </html>
    """
    return html_content

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """Handle WebSocket connections"""
    await websocket.accept()
    connected_clients.add(websocket)
    logger.info(f"Client connected. Total clients: {len(connected_clients)}")
    
    try:
        while True:
            # Just wait for disconnection
            await websocket.receive_text()
    except WebSocketDisconnect:
        connected_clients.discard(websocket)
        logger.info(f"Client disconnected. Total clients: {len(connected_clients)}")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        connected_clients.discard(websocket)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("debug_app:app", host="0.0.0.0", port=7861, reload=False)