client/client.js

// Matrix-Game V2 WebSocket Client

// WebSocket connection
let socket = null;
let userId = null;
let isStreaming = false;
let lastFrameTime = 0;
let frameCount = 0;
let fpsUpdateInterval = null;

// Session state
let sessionInfo = {
    maxFrames: 0,
    sessionDuration: 0,
    currentFrame: 0,
    remainingFrames: 0,
    elapsedTime: 0,
    remainingTime: 0,
    sessionStartTime: null,
    sessionEnded: false
};

// DOM Elements
const connectBtn = document.getElementById('connect-btn');
const startStreamBtn = document.getElementById('start-stream-btn');
const stopStreamBtn = document.getElementById('stop-stream-btn');
const resetSessionBtn = document.getElementById('reset-session-btn');
const sceneSelect = document.getElementById('scene-select');
const gameCanvas = document.getElementById('game-canvas');
const connectionLog = document.getElementById('connection-log');
const mousePosition = document.getElementById('mouse-position');
const fpsCounter = document.getElementById('fps-counter');
const mouseTrackingArea = document.getElementById('mouse-tracking-area');

// Session progress elements
const sessionProgress = document.getElementById('session-progress');
const progressFill = document.getElementById('progress-fill');
const currentFrameEl = document.getElementById('current-frame');
const maxFramesEl = document.getElementById('max-frames');
const elapsedTimeEl = document.getElementById('elapsed-time');
const sessionDurationEl = document.getElementById('session-duration');
const remainingTimeEl = document.getElementById('remaining-time');

// Pointer Lock API support check
const pointerLockSupported = 'pointerLockElement' in document || 
                            'mozPointerLockElement' in document || 
                            'webkitPointerLockElement' in document;

// Keyboard DOM elements
const keyElements = {
    'w': document.getElementById('key-w'),
    'a': document.getElementById('key-a'),
    's': document.getElementById('key-s'),
    'd': document.getElementById('key-d'),
    'space': document.getElementById('key-space'),
    'shift': document.getElementById('key-shift')
};

// Key mapping to action names
const keyToAction = {
    'w': 'forward',
    'arrowup': 'forward',
    'a': 'left',
    'arrowleft': 'left',
    's': 'back',
    'arrowdown': 'back',
    'd': 'right',
    'arrowright': 'right',
    ' ': 'jump',
    'shift': 'attack'
};

// Key state tracking
const keyState = {
    'forward': false,
    'back': false,
    'left': false,
    'right': false,
    'jump': false,
    'attack': false
};

// Mouse state
const mouseState = {
    x: 0,
    y: 0,
    captured: false
};

// Test server connectivity before establishing WebSocket
async function testServerConnectivity() {
    try {
        // Get base path by extracting path from the script tag's src attribute
        let basePath = '';
        const scriptTags = document.getElementsByTagName('script');
        for (const script of scriptTags) {
            if (script.src.includes('client.js')) {
                const url = new URL(script.src);
                basePath = url.pathname.replace('/assets/client.js', '');
                break;
            }
        }
        
        // Try to fetch the debug endpoint to see if the server is accessible
        const response = await fetch(`${window.location.protocol}//${window.location.host}${basePath}/api/debug`);
        if (!response.ok) {
            throw new Error(`Server returned ${response.status}`);
        }
        
        const debugInfo = await response.json();
        logMessage(`Matrix-Game V2 server connected! Server time: ${new Date(debugInfo.server_time * 1000).toLocaleTimeString()}`);
        
        // Check engine status in debug info
        if (debugInfo.server_info && debugInfo.server_info.engine_status) {
            if (debugInfo.server_info.engine_status === 'failed') {
                logMessage(`⚠️  Warning: Engine status is '${debugInfo.server_info.engine_status}'`);
            } else {
                logMessage(`✅ Engine status: ${debugInfo.server_info.engine_status}`);
            }
        }
        
        // Log available routes from server
        if (debugInfo.all_routes && debugInfo.all_routes.length > 0) {
            logMessage(`Available routes: ${debugInfo.all_routes.join(', ')}`);
        }
        
        // Return the debug info for connection setup
        return debugInfo;
    } catch (error) {
        logMessage(`Server connection test failed: ${error.message}`);
        return null;
    }
}

// Connect to WebSocket server
async function connectWebSocket() {
    // First test connectivity to the server
    logMessage('Testing server connectivity...');
    const debugInfo = await testServerConnectivity();
    
    // Use secure WebSocket (wss://) if the page is loaded over HTTPS
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    
    // Get base path by extracting path from the script tag's src attribute
    let basePath = '';
    if (debugInfo && debugInfo.base_path) {
        // Use base path from server if available
        basePath = debugInfo.base_path;
        logMessage(`Using server-provided base path: ${basePath}`);
    } else {
        const scriptTags = document.getElementsByTagName('script');
        for (const script of scriptTags) {
            if (script.src.includes('client.js')) {
                const url = new URL(script.src);
                basePath = url.pathname.replace('/assets/client.js', '');
                break;
            }
        }
    }
    
    // Try both with and without base path for WebSocket connection
    let serverUrl = `${protocol}//${window.location.hostname}${window.location.port ? ':' + window.location.port : ''}${basePath}/ws`;
    logMessage(`Attempting to connect to WebSocket at ${serverUrl}...`);
    
    // For Hugging Face Spaces, try the direct /ws path if the base path doesn't work
    const fallbackUrl = `${protocol}//${window.location.hostname}${window.location.port ? ':' + window.location.port : ''}/ws`;
    
    try {
        socket = new WebSocket(serverUrl);
        setupWebSocketHandlers();
        
        // Set a timeout to try the fallback URL if the first one doesn't connect
        setTimeout(() => {
            if (socket.readyState !== WebSocket.OPEN && socket.readyState !== WebSocket.CONNECTING) {
                logMessage(`Connection to ${serverUrl} failed. Trying fallback URL: ${fallbackUrl}`);
                socket = new WebSocket(fallbackUrl);
                setupWebSocketHandlers();
            }
        }, 3000);
    } catch (error) {
        logMessage(`Error connecting to WebSocket: ${error.message}`);
        resetUI();
    }
}

// Set up WebSocket event handlers
function setupWebSocketHandlers() {
    socket.onopen = () => {
        logMessage('WebSocket connection established');
        connectBtn.textContent = 'Disconnect';
        startStreamBtn.disabled = false;
        stopStreamBtn.disabled = true;
        resetSessionBtn.disabled = false;
        sceneSelect.disabled = false;
    };
    
    socket.onmessage = (event) => {
        const message = JSON.parse(event.data);
        
        switch (message.action) {
            case 'welcome':
                userId = message.userId;
                logMessage(`Welcome to Matrix-Game V2! User ID: ${userId}`);
                
                // Update scene options if server provides them
                if (message.scenes && Array.isArray(message.scenes)) {
                    sceneSelect.innerHTML = '';
                    
                    // Add V2 modes first
                    const v2Modes = ['universal', 'gta_drive', 'temple_run'];
                    const modeNames = {
                        'universal': 'Universal Mode',
                        'gta_drive': 'GTA Drive Mode', 
                        'temple_run': 'Temple Run Mode'
                    };
                    
                    message.scenes.forEach(scene => {
                        const option = document.createElement('option');
                        option.value = scene;
                        
                        // Use friendly names for V2 modes
                        if (modeNames[scene]) {
                            option.textContent = modeNames[scene];
                        } else {
                            // Legacy scenes marked as demo
                            option.textContent = scene.charAt(0).toUpperCase() + scene.slice(1) + ' (Demo)';
                        }
                        
                        // Group V2 modes at the top
                        if (v2Modes.includes(scene)) {
                            sceneSelect.insertBefore(option, sceneSelect.firstChild);
                        } else {
                            sceneSelect.appendChild(option);
                        }
                    });
                    
                    // Default to universal mode if available
                    if (message.scenes.includes('universal')) {
                        sceneSelect.value = 'universal';
                    }
                }
                break;
                
            case 'frame':
                // Process incoming frame
                processFrame(message);
                break;
                
            case 'start_stream':
                if (message.success) {
                    isStreaming = true;
                    sessionInfo.sessionEnded = false;
                    startStreamBtn.disabled = true;
                    stopStreamBtn.disabled = false;
                    logMessage(`Streaming started: ${message.message}`);
                    
                    // Update session info if provided
                    if (message.sessionInfo) {
                        updateSessionProgress(message.sessionInfo);
                    }
                    
                    // Start FPS counter
                    startFpsCounter();
                } else {
                    logMessage(`Error starting stream: ${message.error}`);
                    
                    // Check if session ended
                    if (message.sessionEnded) {
                        sessionInfo.sessionEnded = true;
                        startStreamBtn.disabled = true;
                        stopStreamBtn.disabled = true;
                        logMessage('Session has ended. Please reset to start a new session.');
                    }
                }
                break;
                
            case 'stop_stream':
                if (message.success) {
                    isStreaming = false;
                    startStreamBtn.disabled = false;
                    stopStreamBtn.disabled = true;
                    logMessage('Streaming stopped');
                    
                    // Stop FPS counter
                    stopFpsCounter();
                } else {
                    logMessage(`Error stopping stream: ${message.error}`);
                }
                break;
                
            case 'pong':
                // Server responded to ping
                break;
                
            case 'change_scene':
                if (message.success) {
                    const modeNames = {
                        'universal': 'Universal Mode',
                        'gta_drive': 'GTA Drive Mode',
                        'temple_run': 'Temple Run Mode'
                    };
                    const displayName = modeNames[message.scene] || message.scene;
                    logMessage(`Mode changed to: ${displayName}`);
                } else {
                    logMessage(`Error changing mode: ${message.error}`);
                }
                break;
                
            case 'frame_error':
                logMessage(`❌ Frame Generation Error: ${message.error}`);
                // Stop streaming if there's a frame error
                if (isStreaming) {
                    isStreaming = false;
                    startStreamBtn.disabled = false;
                    stopStreamBtn.disabled = true;
                    startStreamBtn.textContent = 'Start Stream';
                    stopFpsCounter();
                }
                break;
                
            case 'session_ended':
                logMessage(`🏁 ${message.message}`);
                isStreaming = false;
                sessionInfo.sessionEnded = true;
                startStreamBtn.disabled = true;
                stopStreamBtn.disabled = true;
                
                // Show session stats if available
                if (message.sessionStats) {
                    const stats = message.sessionStats;
                    logMessage(`Session Stats: ${stats.totalFrames}/${stats.maxFrames} frames, ${formatTime(stats.elapsedTime)} elapsed`);
                }
                
                // Stop FPS counter
                stopFpsCounter();
                break;
                
            case 'reset_session':
                if (message.success) {
                    logMessage(`🔄 ${message.message}`);
                    
                    // Reset session state
                    sessionInfo = {
                        maxFrames: 0,
                        sessionDuration: 0,
                        currentFrame: 0,
                        remainingFrames: 0,
                        elapsedTime: 0,
                        remainingTime: 0,
                        sessionStartTime: null,
                        sessionEnded: false
                    };
                    
                    // Update UI
                    isStreaming = false;
                    startStreamBtn.disabled = false;
                    stopStreamBtn.disabled = true;
                    sessionProgress.style.display = 'none';
                    
                    // Update session info if provided
                    if (message.sessionInfo) {
                        updateSessionProgress(message.sessionInfo);
                    }
                    
                    stopFpsCounter();
                } else {
                    logMessage(`Error resetting session: ${message.error}`);
                }
                break;
                
            default:
                logMessage(`Received message: ${JSON.stringify(message)}`);
        }
    };
    
    socket.onclose = (event) => {
        logMessage(`WebSocket connection closed (code: ${event.code}, reason: ${event.reason || 'none given'})`);
        resetUI();
    };
    
    socket.onerror = (error) => {
        logMessage(`WebSocket error. This is often caused by CORS issues or the server being inaccessible.`);
        console.error('WebSocket error:', error);
        resetUI();
    };
}

// Disconnect from WebSocket server
function disconnectWebSocket() {
    if (socket && socket.readyState === WebSocket.OPEN) {
        // Stop streaming if active
        if (isStreaming) {
            sendStopStream();
        }
        
        // Close the socket
        socket.close();
        logMessage('Disconnected from server');
    }
}

// Start streaming frames
function sendStartStream() {
    if (socket && socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({
            action: 'start_stream',
            requestId: generateRequestId(),
            fps: 16  // Default FPS
        }));
    }
}

// Stop streaming frames
function sendStopStream() {
    if (socket && socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({
            action: 'stop_stream',
            requestId: generateRequestId()
        }));
    }
}

// Send keyboard input to server
function sendKeyboardInput(key, pressed) {
    if (socket && socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({
            action: 'keyboard_input',
            requestId: generateRequestId(),
            key: key,
            pressed: pressed
        }));
    }
}

// Send mouse input to server
function sendMouseInput(x, y) {
    if (socket && socket.readyState === WebSocket.OPEN && isStreaming) {
        socket.send(JSON.stringify({
            action: 'mouse_input',
            requestId: generateRequestId(),
            x: x,
            y: y
        }));
    }
}

// Change scene
function sendChangeScene(scene) {
    if (socket && socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({
            action: 'change_scene',
            requestId: generateRequestId(),
            scene: scene
        }));
    }
}

// Reset session
function sendResetSession() {
    if (socket && socket.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({
            action: 'reset_session',
            requestId: generateRequestId()
        }));
    }
}

// Update session progress display
function updateSessionProgress(progress) {
    if (progress) {
        sessionInfo = { ...sessionInfo, ...progress };
        
        // Show progress UI if session is active
        if (sessionInfo.maxFrames > 0) {
            sessionProgress.style.display = 'block';
            
            // Update progress bar
            const progressPercent = (sessionInfo.currentFrame / sessionInfo.maxFrames) * 100;
            progressFill.style.width = `${progressPercent}%`;
            
            // Update text displays
            currentFrameEl.textContent = sessionInfo.currentFrame;
            maxFramesEl.textContent = sessionInfo.maxFrames;
            elapsedTimeEl.textContent = formatTime(sessionInfo.elapsedTime);
            sessionDurationEl.textContent = formatTime(sessionInfo.sessionDuration);
            remainingTimeEl.textContent = formatTime(sessionInfo.remainingTime);
        }
    }
}

// Format time in MM:SS format
function formatTime(seconds) {
    if (isNaN(seconds) || seconds < 0) return '0:00';
    const minutes = Math.floor(seconds / 60);
    const secs = Math.floor(seconds % 60);
    return `${minutes}:${secs.toString().padStart(2, '0')}`;
}

// Process incoming frame
function processFrame(message) {
    // Update FPS calculation
    const now = performance.now();
    if (lastFrameTime > 0) {
        frameCount++;
    }
    lastFrameTime = now;
    
    // Update the canvas with the new frame
    if (message.frameData) {
        gameCanvas.src = `data:image/jpeg;base64,${message.frameData}`;
    }
    
    // Update session progress if available
    if (message.sessionProgress) {
        updateSessionProgress(message.sessionProgress);
    }
}

// Generate a random request ID
function generateRequestId() {
    return Math.random().toString(36).substring(2, 15);
}

// Log message to the connection info panel
function logMessage(message) {
    const logEntry = document.createElement('div');
    logEntry.className = 'log-entry';
    
    const timestamp = new Date().toLocaleTimeString();
    logEntry.textContent = `[${timestamp}] ${message}`;
    
    connectionLog.appendChild(logEntry);
    connectionLog.scrollTop = connectionLog.scrollHeight;
    
    // Limit number of log entries
    while (connectionLog.children.length > 100) {
        connectionLog.removeChild(connectionLog.firstChild);
    }
}

// Start FPS counter updates
function startFpsCounter() {
    frameCount = 0;
    lastFrameTime = 0;
    
    // Update FPS display every second
    fpsUpdateInterval = setInterval(() => {
        fpsCounter.textContent = `FPS: ${frameCount}`;
        frameCount = 0;
    }, 1000);
}

// Stop FPS counter updates
function stopFpsCounter() {
    if (fpsUpdateInterval) {
        clearInterval(fpsUpdateInterval);
        fpsUpdateInterval = null;
    }
    fpsCounter.textContent = 'FPS: 0';
}

// Reset UI to initial state
function resetUI() {
    connectBtn.textContent = 'Connect';
    startStreamBtn.disabled = true;
    stopStreamBtn.disabled = true;
    resetSessionBtn.disabled = true;
    sceneSelect.disabled = true;
    
    // Hide session progress
    sessionProgress.style.display = 'none';
    
    // Reset key indicators
    for (const key in keyElements) {
        keyElements[key].classList.remove('active');
    }
    
    // Stop FPS counter
    stopFpsCounter();
    
    // Reset streaming state
    isStreaming = false;
    
    // Reset session state
    sessionInfo = {
        maxFrames: 0,
        sessionDuration: 0,
        currentFrame: 0,
        remainingFrames: 0,
        elapsedTime: 0,
        remainingTime: 0,
        sessionStartTime: null,
        sessionEnded: false
    };
}

// Event Listeners
connectBtn.addEventListener('click', () => {
    if (socket && socket.readyState === WebSocket.OPEN) {
        disconnectWebSocket();
    } else {
        connectWebSocket();
    }
});

startStreamBtn.addEventListener('click', sendStartStream);
stopStreamBtn.addEventListener('click', sendStopStream);
resetSessionBtn.addEventListener('click', sendResetSession);

sceneSelect.addEventListener('change', () => {
    sendChangeScene(sceneSelect.value);
});

// Keyboard event listeners
document.addEventListener('keydown', (event) => {
    const key = event.key.toLowerCase();
    
    // Map key to action
    let action = keyToAction[key];
    if (!action && key === ' ') {
        action = keyToAction[' '];  // Handle spacebar
    }
    
    if (action && !keyState[action]) {
        keyState[action] = true;
        
        // Update visual indicator
        const keyElement = keyElements[key] || 
                          (key === ' ' ? keyElements['space'] : null) ||
                          (key === 'shift' ? keyElements['shift'] : null);
        
        if (keyElement) {
            keyElement.classList.add('active');
        }
        
        // Send to server
        sendKeyboardInput(action, true);
    }
    
    // Prevent default actions for game controls
    if (Object.keys(keyToAction).includes(key) || key === ' ') {
        event.preventDefault();
    }
});

document.addEventListener('keyup', (event) => {
    const key = event.key.toLowerCase();
    
    // Map key to action
    let action = keyToAction[key];
    if (!action && key === ' ') {
        action = keyToAction[' '];  // Handle spacebar
    }
    
    if (action && keyState[action]) {
        keyState[action] = false;
        
        // Update visual indicator
        const keyElement = keyElements[key] || 
                          (key === ' ' ? keyElements['space'] : null) ||
                          (key === 'shift' ? keyElements['shift'] : null);
        
        if (keyElement) {
            keyElement.classList.remove('active');
        }
        
        // Send to server
        sendKeyboardInput(action, false);
    }
});

// Mouse capture functions
function requestPointerLock() {
    if (!mouseState.captured && pointerLockSupported) {
        mouseTrackingArea.requestPointerLock = mouseTrackingArea.requestPointerLock ||
                                            mouseTrackingArea.mozRequestPointerLock ||
                                            mouseTrackingArea.webkitRequestPointerLock;
        mouseTrackingArea.requestPointerLock();
        logMessage('Mouse captured. Press ESC to release.');
    }
}

function exitPointerLock() {
    if (mouseState.captured) {
        document.exitPointerLock = document.exitPointerLock ||
                                 document.mozExitPointerLock ||
                                 document.webkitExitPointerLock;
        document.exitPointerLock();
        logMessage('Mouse released.');
    }
}

// Handle pointer lock change events
document.addEventListener('pointerlockchange', pointerLockChangeHandler);
document.addEventListener('mozpointerlockchange', pointerLockChangeHandler);
document.addEventListener('webkitpointerlockchange', pointerLockChangeHandler);

function pointerLockChangeHandler() {
    if (document.pointerLockElement === mouseTrackingArea ||
        document.mozPointerLockElement === mouseTrackingArea ||
        document.webkitPointerLockElement === mouseTrackingArea) {
        // Pointer is locked, enable mouse movement tracking
        mouseState.captured = true;
        document.addEventListener('mousemove', handleMouseMovement);
    } else {
        // Pointer is unlocked, disable mouse movement tracking
        mouseState.captured = false;
        document.removeEventListener('mousemove', handleMouseMovement);
        // Reset mouse state
        mouseState.x = 0;
        mouseState.y = 0;
        mousePosition.textContent = `Mouse: ${mouseState.x.toFixed(2)}, ${mouseState.y.toFixed(2)}`;
        throttledSendMouseInput();
    }
}

// Mouse tracking with pointer lock
function handleMouseMovement(event) {
    if (mouseState.captured) {
        // Use movement for mouse look when captured
        const sensitivity = 0.005; // Adjust sensitivity
        mouseState.x += event.movementX * sensitivity;
        mouseState.y -= event.movementY * sensitivity; // Invert Y for intuitive camera control
        
        // Clamp values
        mouseState.x = Math.max(-1, Math.min(1, mouseState.x));
        mouseState.y = Math.max(-1, Math.min(1, mouseState.y));
        
        // Update display
        mousePosition.textContent = `Mouse: ${mouseState.x.toFixed(2)}, ${mouseState.y.toFixed(2)}`;
        
        // Send to server (throttled)
        throttledSendMouseInput();
    }
}

// Mouse click to capture
mouseTrackingArea.addEventListener('click', () => {
    if (!mouseState.captured && isStreaming) {
        requestPointerLock();
    }
});

// Standard mouse tracking for when pointer is not locked
mouseTrackingArea.addEventListener('mousemove', (event) => {
    if (!mouseState.captured) {
        // Calculate normalized coordinates relative to the center of the tracking area
        const rect = mouseTrackingArea.getBoundingClientRect();
        const centerX = rect.width / 2;
        const centerY = rect.height / 2;
        
        // Calculate relative position from center (-1 to 1)
        const relX = (event.clientX - rect.left - centerX) / centerX;
        const relY = (event.clientY - rect.top - centerY) / centerY;
        
        // Scale down for smoother movement (similar to conditions.py)
        const scaleFactor = 0.05;
        mouseState.x = relX * scaleFactor;
        mouseState.y = -relY * scaleFactor;  // Invert Y for intuitive camera control
        
        // Update display
        mousePosition.textContent = `Mouse: ${mouseState.x.toFixed(2)}, ${mouseState.y.toFixed(2)}`;
        
        // Send to server (throttled)
        throttledSendMouseInput();
    }
});

// Throttle mouse movement to avoid flooding the server
const throttledSendMouseInput = (() => {
    let lastSentTime = 0;
    const interval = 50;  // milliseconds
    
    return () => {
        const now = performance.now();
        if (now - lastSentTime >= interval) {
            sendMouseInput(mouseState.x, mouseState.y);
            lastSentTime = now;
        }
    };
})();

// Toggle panel collapse/expand
function togglePanel(panelId) {
    const panel = document.getElementById(panelId);
    const button = panel.querySelector('.toggle-button');
    
    if (panel.classList.contains('collapsed')) {
        // Expand the panel
        panel.classList.remove('collapsed');
        button.textContent = '−'; // Minus sign
    } else {
        // Collapse the panel
        panel.classList.add('collapsed');
        button.textContent = '+'; // Plus sign
    }
}

// Initialize the UI
resetUI();

// Make panel headers clickable
document.querySelectorAll('.panel-header').forEach(header => {
    header.addEventListener('click', () => {
        const panelId = header.parentElement.id;
        togglePanel(panelId);
    });
});