// Initialize variables
let isVirtual = false;

// DOM elements
const videoContainer = document.getElementById('videoContainer');
const videoElement = document.getElementById('videoElement');
const canvasElement = document.getElementById('backgroundCanvas');
const backgroundImage = document.getElementById('yourBackgroundImage');
const ctx = canvasElement.getContext('2d');
const startButton = document.getElementById('startButton');
const stopButton = document.getElementById('stopButton');
const errorText = document.getElementById('errorText');

 // MobileNetV1 or ResNet50
const BodyPixModel = 'MobileNetV1';

async function startWebCamStream() {
    try {
        // Start the webcam stream
        const stream = await navigator.mediaDevices.getUserMedia({ video: true });
        videoElement.srcObject = stream;

        // Wait for the video to play
        await videoElement.play();
    } catch (error) {
        displayError(error)
    }
}

// Function to start the virtual background
async function startVirtualBackground() {
    try {
        // Set canvas dimensions to match video dimensions
        canvasElement.width = videoElement.videoWidth;
        canvasElement.height = videoElement.videoHeight;

        // Load the BodyPix model: 
        let net = null;

        switch (BodyPixModel) {
            case 'MobileNetV1':
                /*
                    This is a lightweight architecture that is suitable for real-time applications and has lower computational requirements. 
                    It provides good performance for most use cases.
                */
                net = await bodyPix.load({
                    architecture: 'MobileNetV1',
                    outputStride: 16, // Output stride (16 or 32). 16 is faster, 32 is more accurate.
                    multiplier: 0.75, // The model's depth multiplier. Options: 0.50, 0.75, or 1.0.
                    quantBytes: 2, // The number of bytes to use for quantization (4 or 2).
                });
                break;
            case 'ResNet50':
                /*
                    This is a deeper and more accurate architecture compared to MobileNetV1. 
                    It may provide better segmentation accuracy, but it requires more computational resources and can be slower.
                */
                net = await bodyPix.load({
                    architecture: 'ResNet50',
                    outputStride: 16, // Output stride (16 or 32). 16 is faster, 32 is more accurate.
                    quantBytes: 4, // The number of bytes to use for quantization (4 or 2).
                });
                break;
            default:
                break;
        }

        // Start the virtual background loop
        isVirtual = true;
        videoElement.hidden = true;
        canvasElement.hidden = false;
        display(canvasElement, 'block');

        // Show the stop button and hide the start button
        startButton.style.display = 'none';
        stopButton.style.display = 'block';

        async function updateCanvas() {
            if (isVirtual) {
                // 1. Segmentation Calculation
                const segmentation = await net.segmentPerson(videoElement, {
                    flipHorizontal: false, // Whether to flip the input video horizontally
                    internalResolution: 'medium', // The resolution for internal processing (options: 'low', 'medium', 'high')
                    segmentationThreshold: 0.7, // Segmentation confidence threshold (0.0 - 1.0)
                    maxDetections: 10, // Maximum number of detections to return
                    scoreThreshold: 0.2, // Confidence score threshold for detections (0.0 - 1.0)
                    nmsRadius: 20, // Non-Maximum Suppression (NMS) radius for de-duplication
                    minKeypointScore: 0.3, // Minimum keypoint detection score (0.0 - 1.0)
                    refineSteps: 10, // Number of refinement steps for segmentation
                });

                // 2. Creating a Background Mask
                const background = { r: 0, g: 0, b: 0, a: 0 };
                const mask = bodyPix.toMask(segmentation, background, { r: 0, g: 0, b: 0, a: 255 });

                if (mask) {
                    ctx.putImageData(mask, 0, 0);
                    ctx.globalCompositeOperation = 'source-in';

                    // 3. Drawing the Background
                    if (backgroundImage.complete) {
                        ctx.drawImage(backgroundImage, 0, 0, canvasElement.width, canvasElement.height);
                    } else {
                        // If the image is not loaded yet, wait for it to load and then draw
                        backgroundImage.onload = () => {
                            ctx.drawImage(backgroundImage, 0, 0, canvasElement.width, canvasElement.height);
                        };
                    }

                    // Draw the mask (segmentation)
                    ctx.globalCompositeOperation = 'destination-over';
                    ctx.drawImage(videoElement, 0, 0, canvasElement.width, canvasElement.height);
                    ctx.globalCompositeOperation = 'source-over';

                    // Add a delay to control the frame rate (adjust as needed) less CPU intensive
                    // await new Promise((resolve) => setTimeout(resolve, 100));

                    // Continue updating the canvas
                    requestAnimationFrame(updateCanvas);
                }
            }
        }
        // Start update canvas
        updateCanvas();
    } catch (error) {
        stopVirtualBackground();
        displayError(error);
    }
}

// Function to stop the virtual background
function stopVirtualBackground() {
    isVirtual = false;
    videoElement.hidden = false;
    canvasElement.hidden = true;
    display(canvasElement, 'none');

    // Hide the stop button and show the start button
    startButton.style.display = 'block';
    stopButton.style.display = 'none';
}

// Helper function to set the display style of an element
function display(element, style) {
    element.style.display = style;
}

// Helper function to display errors
function displayError(error){
    console.error(error);
    // Display error message in the <p> element
    errorText.textContent = 'An error occurred: ' + error.message;
}

// Add click event listeners to the buttons
startButton.addEventListener('click', startVirtualBackground);
stopButton.addEventListener('click', stopVirtualBackground);

// Start video stream
startWebCamStream();