static/js/websocket_trainer.js

document.addEventListener('DOMContentLoaded', () => {
    const videoElement = document.getElementById('userVideo');
    const canvasElement = document.createElement('canvas'); // Offscreen canvas
    const context = canvasElement.getContext('2d');

    const startButton = document.getElementById('startTrainerBtn');
    const stopButton = document.getElementById('stopTrainerBtn');
    const exerciseTypeSelect = document.getElementById('exerciseType');

    // UI Displays for feedback
    const repsDisplay = document.getElementById('repsDisplay');
    const stageDisplay = document.getElementById('stageDisplay');
    const feedbackDisplay = document.getElementById('feedbackDisplay');
    const angleDisplay = document.getElementById('angleDisplay');

    const repsLeftHcDisplay = document.getElementById('repsLeftHc');
    const repsRightHcDisplay = document.getElementById('repsRightHc');
    const stageLeftHcDisplay = document.getElementById('stageLeftHc');
    const stageRightHcDisplay = document.getElementById('stageRightHc');
    const feedbackLeftHcDisplay = document.getElementById('feedbackLeftHc');
    const feedbackRightHcDisplay = document.getElementById('feedbackRightHc');
    const angleLeftCurlHcDisplay = document.getElementById('angleLeftCurlHc');
    const angleRightCurlHcDisplay = document.getElementById('angleRightCurlHc');

    const squatPushupUIDiv = document.getElementById('squatPushupUI');
    const hammerCurlUIDiv = document.getElementById('hammerCurlUI');

    const apiStatusDisplay = document.getElementById('apiStatus');
    const sessionIdDisplay = document.getElementById('sessionIdDisplay');

    let socket = null;
    let currentExerciseType = '';
    let isSessionActive = false;
    let animationFrameId = null;
    let stream = null;

    const logger = {
        log: (message) => console.log(`[Trainer Log] ${new Date().toISOString()}: ${message}`),
        error: (message) => console.error(`[Trainer Error] ${new Date().toISOString()}: ${message}`),
        info: (message) => console.info(`[Trainer Info] ${new Date().toISOString()}: ${message}`)
    };

    logger.log("WebSocket trainer script loaded. Waiting for DOM and user interaction.");

    function setupUIForExercise(exerciseType) {
        logger.info(`Setting up UI for exercise: ${exerciseType}`);
        if (exerciseType === 'hammer_curl') {
            if (squatPushupUIDiv) squatPushupUIDiv.style.display = 'none';
            if (hammerCurlUIDiv) hammerCurlUIDiv.style.display = 'block';
        } else {
            if (squatPushupUIDiv) squatPushupUIDiv.style.display = 'block';
            if (hammerCurlUIDiv) hammerCurlUIDiv.style.display = 'none';
        }
    }

    function updateGenericUI(data) {
        if (repsDisplay) repsDisplay.textContent = data.counter !== undefined ? data.counter : 'N/A';
        if (stageDisplay) stageDisplay.textContent = data.stage || 'N/A';
        if (feedbackDisplay) feedbackDisplay.textContent = data.feedback || 'N/A';

        let angleText = 'Angle: -';
        if (data.angle_left !== undefined && data.angle_right !== undefined) {
            angleText = `Angles L: ${Math.round(data.angle_left)}, R: ${Math.round(data.angle_right)}`;
        } else if (data.angle_body_left !== undefined && data.angle_body_right !== undefined) {
            angleText = `Body Angles L: ${Math.round(data.angle_body_left)}, R: ${Math.round(data.angle_body_right)}`;
        }
        if (angleDisplay) angleDisplay.textContent = angleText;
    }

    function updateHammerCurlUI(data) {
        if (repsLeftHcDisplay) repsLeftHcDisplay.textContent = data.counter_left !== undefined ? data.counter_left : 'N/A';
        if (repsRightHcDisplay) repsRightHcDisplay.textContent = data.counter_right !== undefined ? data.counter_right : 'N/A';
        if (stageLeftHcDisplay) stageLeftHcDisplay.textContent = data.stage_left || 'N/A';
        if (stageRightHcDisplay) stageRightHcDisplay.textContent = data.stage_right || 'N/A';
        if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = data.feedback_left || 'N/A';
        if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = data.feedback_right || 'N/A';

        let curlAngleText = 'Curl Angles: -';
        if (data.angle_left_curl !== undefined && data.angle_right_curl !== undefined) {
            curlAngleText = `Curl Angles L: ${Math.round(data.angle_left_curl)}, R: ${Math.round(data.angle_right_curl)}`;
        }
        // Assuming you want to display curl angles in the general angleDisplay or have a specific one
        if (angleDisplay) angleDisplay.textContent = curlAngleText;
    }

    function clearUIFeedback() {
        logger.info("Clearing UI feedback elements.");
        if (repsDisplay) repsDisplay.textContent = '0';
        if (feedbackDisplay) feedbackDisplay.textContent = 'N/A';
        if (stageDisplay) stageDisplay.textContent = 'N/A';
        if (angleDisplay) angleDisplay.textContent = 'Angle: -';

        if (repsLeftHcDisplay) repsLeftHcDisplay.textContent = '0';
        if (repsRightHcDisplay) repsRightHcDisplay.textContent = '0';
        if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = 'N/A';
        if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = 'N/A';
        if (stageLeftHcDisplay) stageLeftHcDisplay.textContent = 'N/A';
        if (stageRightHcDisplay) stageRightHcDisplay.textContent = 'N/A';
        if (angleLeftCurlHcDisplay) angleLeftCurlHcDisplay.textContent = '0';
        if (angleRightCurlHcDisplay) angleRightCurlHcDisplay.textContent = '0';

        if (apiStatusDisplay) apiStatusDisplay.textContent = "Idle";
        if (sessionIdDisplay) sessionIdDisplay.textContent = "-";
    }

    function initializeSocket() {
        if (socket && socket.connected) {
            logger.info("Socket already connected.");
            return;
        }
        // Ensure the server URL is correct, especially if not running on the same host/port
        // For Hugging Face Spaces, it should connect to the same origin, so io() is fine.
        socket = io();

        socket.on('connect', () => {
            logger.info(`Connected to server with SID: ${socket.id}`);
            if (apiStatusDisplay) apiStatusDisplay.textContent = "Connected to server.";
        });

        socket.on('disconnect', (reason) => {
            logger.info(`Disconnected from server: ${reason}`);
            if (apiStatusDisplay) apiStatusDisplay.textContent = "Disconnected. Refresh if needed.";
            stopTrainer(false);
        });

        socket.on('connect_error', (error) => {
            logger.error(`Connection Error: ${error.message}`);
            if (apiStatusDisplay) apiStatusDisplay.textContent = "Connection Error.";
        });

        socket.on('connection_ack', (data) => {
            logger.info(`Server Acknowledged Connection: ${JSON.stringify(data)}`);
            if (sessionIdDisplay && data.sid) sessionIdDisplay.textContent = data.sid;
        });

        socket.on('session_started', (data) => {
            logger.info(`Session started on server: ${data.session_id} for exercise: ${data.exercise_type}`);
            isSessionActive = true;
            // currentExerciseType should have been set by startButton click
            if (apiStatusDisplay) apiStatusDisplay.textContent = `Session for ${currentExerciseType} started.`;

            // Check if video is already playing and ready
            if (videoElement && videoElement.srcObject && !videoElement.paused && videoElement.readyState >= videoElement.HAVE_FUTURE_DATA) {
                logger.info("Video ready and session started, initiating frame loop.");
                startFrameLoop();
            } else {
                logger.info("Session started, but video not yet playing/ready. Frame loop will start via video 'onplaying' or 'oncanplaythrough' event.");
            }
        });

        socket.on('session_error', (data) => {
            logger.error(`Session error: ${data.error}`);
            if (apiStatusDisplay) apiStatusDisplay.textContent = `Session Error: ${data.error}`;
            alert(`Session Error: ${data.error}`);
            stopTrainer(false);
        });

        socket.on('exercise_update', (data) => {
            if (!isSessionActive) return;

            if (data.success) {
                if (data.landmarks_detected) {
                    // logger.log('Exercise Update:', data.data);
                    if (currentExerciseType === 'hammer_curl') {
                        updateHammerCurlUI(data.data);
                    } else {
                        updateGenericUI(data.data);
                    }
                } else {
                    if (currentExerciseType === 'hammer_curl') {
                        if (feedbackLeftHcDisplay) feedbackLeftHcDisplay.textContent = `Feedback: ${data.message || 'No landmarks detected.'}`;
                        if (feedbackRightHcDisplay) feedbackRightHcDisplay.textContent = '';
                    } else {
                        if (feedbackDisplay) feedbackDisplay.textContent = `Feedback: ${data.message || 'No landmarks detected.'}`;
                    }
                }
            } else {
                logger.error(`Exercise update indicated failure: ${data.message || 'Unknown error'}`);
                if (feedbackDisplay && currentExerciseType !== 'hammer_curl') feedbackDisplay.textContent = `Feedback: Error - ${data.message || 'Processing error'}`;
                else if (feedbackLeftHcDisplay && currentExerciseType === 'hammer_curl') feedbackLeftHcDisplay.textContent = `Feedback: Error - ${data.message || 'Processing error'}`;
            }
        });

        socket.on('frame_error', (data) => {
            logger.error(`Frame processing error from server: ${data.error}`);
            if (feedbackDisplay && currentExerciseType !== 'hammer_curl') feedbackDisplay.textContent = `Error: ${data.error}`;
            else if (feedbackLeftHcDisplay && currentExerciseType === 'hammer_curl') feedbackLeftHcDisplay.textContent = `Error: ${data.error}`;
        });
    }

    function startFrameLoop() {
        logger.info(`Attempting to start frame loop. isSessionActive: ${isSessionActive}, video.paused: ${videoElement.paused}, video.ended: ${videoElement.ended}, video.readyState: ${videoElement.readyState}`);
        if (isSessionActive && videoElement && videoElement.srcObject && !videoElement.paused && !videoElement.ended && videoElement.readyState >= videoElement.HAVE_ENOUGH_DATA) { //HAVE_ENOUGH_DATA is 4
            logger.info("Conditions met, starting sendFrameLoop via requestAnimationFrame.");
            if (animationFrameId) {
                cancelAnimationFrame(animationFrameId);
            }
            sendFrameLoop();
        } else {
            logger.warn(`Conditions not met to start frame loop. isSessionActive: ${isSessionActive}, srcObject: ${!!videoElement.srcObject}, paused: ${videoElement.paused}, ended: ${videoElement.ended}, readyState: ${videoElement.readyState}. Will try again if video becomes ready.`);
        }
    }

    async function sendFrameLoop() {
        if (!isSessionActive || !videoElement.srcObject || videoElement.paused || videoElement.ended || videoElement.readyState < videoElement.HAVE_FUTURE_DATA) { // HAVE_FUTURE_DATA is 3
            logger.info(`[sendFrameLoop] Stopping. Conditions: isSessionActive=${isSessionActive}, srcObject=${!!videoElement.srcObject}, paused=${videoElement.paused}, ended=${videoElement.ended}, readyState=${videoElement.readyState}`);
            if (animationFrameId) {
                cancelAnimationFrame(animationFrameId);
                animationFrameId = null;
            }
            // If session is supposed to be active but video state is bad, consider stopping the session.
            if (isSessionActive && (videoElement.paused || videoElement.ended || videoElement.readyState < videoElement.HAVE_FUTURE_DATA)) {
                logger.warn("[sendFrameLoop] Video stream issue detected, stopping session from client-side.");
                // stopTrainer(true); // This might cause a loop if called from here. Better to just stop sending.
            }
            return;
        }

        try {
            if (canvasElement.width > 0 && canvasElement.height > 0) {
                context.drawImage(videoElement, 0, 0, canvasElement.width, canvasElement.height);
                const imageDataBase64 = canvasElement.toDataURL('image/jpeg', 0.7).split(',')[1];

                if (socket && socket.connected) {
                    // console.log(`[sendFrameLoop] Sending frame for exercise: ${currentExerciseType}`); // Can be too verbose
                    socket.emit('process_frame', {
                        image: imageDataBase64,
                        exercise_type: currentExerciseType,
                        frame_width: canvasElement.width,
                        frame_height: canvasElement.height
                    });
                } else {
                    logger.error("[sendFrameLoop] Socket not connected. Cannot send frame.");
                    stopTrainer(false);
                }
            } else {
                // This might happen if video metadata hasn't loaded yet, or if the video dimensions are truly zero.
                logger.warn(`[sendFrameLoop] Canvas dimensions are zero (w:${canvasElement.width}, h:${canvasElement.height}) or video not ready, skipping frame draw/send.`);
            }
        } catch (error) {
            logger.error("Error in sendFrameLoop (drawing or emitting):", error);
            stopTrainer(false);
        }

        if (isSessionActive) {
            animationFrameId = requestAnimationFrame(sendFrameLoop);
        }
    }

    async function startTrainer() {
        if (isSessionActive) {
            logger.info("Session already active. Please stop current session first.");
            return;
        }

        currentExerciseType = exerciseTypeSelect.value;
        setupUIForExercise(currentExerciseType);
        logger.info(`Attempting to start exercise: ${currentExerciseType}`);
        if (apiStatusDisplay) apiStatusDisplay.textContent = "Initializing...";
        if (startButton) startButton.disabled = true;
        if (stopButton) stopButton.disabled = false;
        if (exerciseTypeSelect) exerciseTypeSelect.disabled = true;

        if (!socket || !socket.connected) {
            logger.info("Socket not connected. Initializing socket connection...");
            initializeSocket();
        }

        // Emit start_exercise_session. The 'session_started' event will then trigger the frame loop.
        // Add a small delay to ensure socket might connect if it was just initialized.
        setTimeout(() => {
            if (socket && socket.connected) {
                logger.info(`Socket connected, emitting start_exercise_session for ${currentExerciseType}`);
                socket.emit('start_exercise_session', { exercise_type: currentExerciseType });
            } else {
                logger.error("Socket not connected after delay. Cannot start session.");
                if (apiStatusDisplay) apiStatusDisplay.textContent = "Error: Cannot connect to server.";
                if (startButton) startButton.disabled = false;
                if (stopButton) stopButton.disabled = true;
                if (exerciseTypeSelect) exerciseTypeSelect.disabled = false;
            }
        }, 500); // 500ms delay to allow socket to connect if just initialized

        try {
            logger.info("Requesting user media...");
            stream = await navigator.mediaDevices.getUserMedia({ video: { width: 640, height: 480 }, audio: false });
            logger.info("User media obtained.");
            videoElement.srcObject = stream;

            videoElement.onloadedmetadata = () => {
                logger.info("Video metadata loaded.");
                canvasElement.width = videoElement.videoWidth;
                canvasElement.height = videoElement.videoHeight;
                logger.info(`Canvas dimensions set to: ${canvasElement.width}x${canvasElement.height}`);
                videoElement.play().catch(playError => {
                    logger.error(`Error playing video: ${playError}`);
                    if (apiStatusDisplay) apiStatusDisplay.textContent = "Error playing video.";
                    alert(`Error playing video: ${playError.message}`);
                    stopTrainer(false);
                });
            };

            videoElement.oncanplay = () => {
                logger.info("Video event: canplay. Video should be ready to play.");
                // Attempt to start frame loop if session is active and video is playable
                if(isSessionActive && videoElement.readyState >= videoElement.HAVE_FUTURE_DATA) {
                    startFrameLoop();
                }
            };

            videoElement.oncanplaythrough = () => {
                logger.info("Video event: canplaythrough. Video has enough data to play without interruption.");
                if(isSessionActive) {
                    startFrameLoop();
                }
            };

            videoElement.onplaying = () => {
                logger.info("Video event: playing. Video has started playing.");
                if(isSessionActive) {
                    startFrameLoop();
                }
            };

            videoElement.onpause = () => {
                logger.info("Video event: pause. Stopping frame loop.");
                isSessionActive = false;
                if (animationFrameId) {
                    cancelAnimationFrame(animationFrameId);
                    animationFrameId = null;
                }
            };

            videoElement.onended = () => {
                logger.info("Video event: ended. Stopping trainer.");
                stopTrainer(true);
            };

        } catch (err) {
            logger.error(`Error accessing webcam: ${err}`);
            if (apiStatusDisplay) apiStatusDisplay.textContent = "Error accessing webcam.";
            alert(`Could not access webcam: ${err.message}`);
            if (startButton) startButton.disabled = false;
            if (stopButton) stopButton.disabled = true;
            if (exerciseTypeSelect) exerciseTypeSelect.disabled = false;
        }
    }

    function stopTrainer(emitToServer = true) {
        logger.info("Stop trainer called.");
        isSessionActive = false;
        if (animationFrameId) {
            cancelAnimationFrame(animationFrameId);
            animationFrameId = null;
            logger.info("Frame loop cancelled.");
        }

        if (stream) {
            stream.getTracks().forEach(track => track.stop());
            stream = null;
            logger.info("Webcam tracks stopped.");
        }
        if (videoElement) {
            videoElement.srcObject = null;
            videoElement.pause(); // Explicitly pause
        }


        if (socket && socket.connected && emitToServer) {
            // No specific 'end_exercise_session' event needed as server handles 'disconnect'
            // but if you want explicit cleanup, you can add one.
            // socket.emit('stop_exercise_session');
            logger.info("Client-side stop initiated. Server will handle via disconnect or if an explicit stop event is implemented.");
        }

        if (apiStatusDisplay) apiStatusDisplay.textContent = "Trainer stopped. Select an exercise to start.";
        if (startButton) startButton.disabled = false;
        if (stopButton) stopButton.disabled = true;
        if (exerciseTypeSelect) exerciseTypeSelect.disabled = false;

        clearUIFeedback();
        if (sessionIdDisplay) sessionIdDisplay.textContent = '-';
    }

    // Initialize Socket.IO connection when the script loads
    if (startButton && stopButton && exerciseTypeSelect) { // Ensure essential controls are present
        initializeSocket();

        startButton.addEventListener('click', startTrainer);
        stopButton.addEventListener('click', () => stopTrainer(true));

        exerciseTypeSelect.addEventListener('change', (event) => {
            if (!isSessionActive) {
                currentExerciseType = event.target.value;
                setupUIForExercise(currentExerciseType);
            }
        });
        // Initial UI setup based on default selected exercise
        setupUIForExercise(exerciseTypeSelect.value);
    } else {
        logger.error("One or more essential UI elements (startButton, stopButton, exerciseTypeSelect) not found.");
    }

    clearUIFeedback();
});