/*
    YOLO VISION
    Browser controller

    Responsibilities:

    1. User selects ONNX model
    2. Read model locally
    3. Send model to Web Worker
    4. Access webcam
    5. Send camera frames to Worker
    6. Receive detections
    7. Draw boxes + segmentation masks
*/


const INPUT_SIZE = 640;


/* =========================================================
   STATE
========================================================= */

let worker = null;

let stream = null;

let running = false;

let frameBusy = false;

let confidenceThreshold = 0.25;

let fps = 0;

let lastFrameTime = 0;


/* =========================================================
   DOM
========================================================= */

const modelFile =
    document.getElementById(
        "modelFile"
    );


const modelStatus =
    document.getElementById(
        "modelStatus"
    );


const modelFileName =
    document.getElementById(
        "modelFileName"
    );


const loadingStatus =
    document.getElementById(
        "loadingStatus"
    );


const progressBar =
    document.getElementById(
        "progressBar"
    );


const progressPercent =
    document.getElementById(
        "progressPercent"
    );


const status =
    document.getElementById(
        "status"
    );


const backend =
    document.getElementById(
        "backend"
    );


const camera =
    document.getElementById(
        "camera"
    );


const output =
    document.getElementById(
        "output"
    );


const outputCtx =
    output.getContext(
        "2d"
    );


const cameraPlaceholder =
    document.getElementById(
        "cameraPlaceholder"
    );


const outputPlaceholder =
    document.getElementById(
        "outputPlaceholder"
    );


const startButton =
    document.getElementById(
        "start"
    );


const stopButton =
    document.getElementById(
        "stop"
    );


const confidence =
    document.getElementById(
        "confidence"
    );


const confidenceText =
    document.getElementById(
        "confidenceText"
    );


const objectsText =
    document.getElementById(
        "objects"
    );


/* =========================================================
   PROGRESS
========================================================= */

function updateProgress(
    percent,
    message
) {

    percent =
        Math.max(
            0,
            Math.min(
                100,
                percent
            )
        );


    progressBar.style.width =
        `${percent}%`;


    progressPercent.textContent =
        `${Math.round(percent)}%`;


    loadingStatus.textContent =
        message;
}


/* =========================================================
   FORMAT FILE SIZE
========================================================= */

function formatBytes(bytes) {

    if (bytes < 1024) {

        return `${bytes} B`;

    }


    if (
        bytes <
        1024 * 1024
    ) {

        return (
            bytes /
            1024
        ).toFixed(1)
        + " KB";

    }


    return (
        bytes /
        1024 /
        1024
    ).toFixed(2)
    + " MB";
}


/* =========================================================
   CREATE WORKER
========================================================= */

function createWorker() {

    if (worker) {

        worker.terminate();

    }


    worker =
        new Worker(
            "worker.js",
            {
                type: "module"
            }
        );


    worker.onmessage =
        handleWorkerMessage;


    worker.onerror =
        function(event) {

            console.error(
                "Worker error:",
                event
            );


            modelStatus.textContent =
                "🔴 Worker error";


            status.textContent =
                "🔴 Worker error";


            updateProgress(
                0,
                "Worker error. Check browser console."
            );

        };

}


/* =========================================================
   MODEL FILE SELECTION
========================================================= */

modelFile.addEventListener(
    "change",
    async function(event) {

        const file =
            event.target.files[0];


        if (!file) {

            return;

        }


        console.log(
            "Selected:",
            file.name
        );


        /*
            Basic validation
        */

        if (
            !file.name
                .toLowerCase()
                .endsWith(".onnx")
        ) {

            alert(
                "Please select an ONNX (.onnx) model."
            );

            return;

        }


        modelFileName.textContent =
            `${file.name} — ${formatBytes(file.size)}`;


        modelStatus.textContent =
            "🟡 Reading model...";


        status.textContent =
            "🟡 Loading YOLO model";


        updateProgress(
            10,
            "Reading selected ONNX model..."
        );


        try {

            /*
                Read model into memory
            */

            const buffer =
                await file.arrayBuffer();


            updateProgress(
                45,
                "Model loaded into browser memory..."
            );


            /*
                Create worker
            */

            createWorker();


            updateProgress(
                65,
                "Sending model to AI worker..."
            );


            /*
                Transfer ArrayBuffer.

                Transfer means the browser does not
                make another copy of the model.
            */

            worker.postMessage(
                {
                    type: "init",

                    modelBytes:
                        buffer,

                    inputSize:
                        INPUT_SIZE

                },
                [
                    buffer
                ]
            );

        }

        catch(error) {

            console.error(
                error
            );


            modelStatus.textContent =
                "🔴 Model loading failed";


            status.textContent =
                "🔴 Model Error";


            updateProgress(
                0,
                "❌ " + error.message
            );

        }

    }
);


/* =========================================================
   WORKER MESSAGE HANDLER
========================================================= */

function handleWorkerMessage(event) {

    const data =
        event.data;


    /*
        Worker setup progress
    */

    if (
        data.type ===
        "loading"
    ) {

        updateProgress(
            data.progress,
            data.message
        );

        return;

    }


    /*
        Model ready
    */

    if (
        data.type ===
        "ready"
    ) {

        updateProgress(
            100,
            "YOLO model ready!"
        );


        modelStatus.textContent =
            "🟢 YOLO11n-Seg Ready";


        status.textContent =
            "🟢 YOLO Ready";


        backend.textContent =
            data.backend;


        startButton.disabled =
            false;


        return;

    }


    /*
        Error
    */

    if (
        data.type ===
        "error"
    ) {

        console.error(
            data.message
        );


        modelStatus.textContent =
            "🔴 Model error";


        status.textContent =
            "🔴 AI Error";


        updateProgress(
            0,
            "❌ " + data.message
        );


        return;

    }


    /*
        Detection result
    */

    if (
        data.type ===
        "result"
    ) {

        frameBusy =
            false;


        drawResult(
            data
        );

    }

}


/* =========================================================
   START CAMERA
========================================================= */

async function startCamera() {

    if (!worker) {

        alert(
            "Please select the YOLO ONNX model first."
        );

        return;

    }


    try {

        stream =
            await navigator
                .mediaDevices
                .getUserMedia(
                    {
                        video: {
                            facingMode:
                                "environment"
                        },

                        audio: false
                    }
                );


        camera.srcObject =
            stream;


        await camera.play();


        running =
            true;


        startButton.disabled =
            true;


        stopButton.disabled =
            false;


        cameraPlaceholder.style.display =
            "none";


        outputPlaceholder.style.display =
            "none";


        status.textContent =
            "🟢 Camera Running";


        processFrame();

    }

    catch(error) {

        console.error(
            error
        );


        status.textContent =
            "🔴 Camera Error";


        alert(
            "Camera access failed:\n" +
            error.message
        );

    }

}


/* =========================================================
   STOP CAMERA
========================================================= */

function stopCamera() {

    running =
        false;


    if (stream) {

        stream
            .getTracks()
            .forEach(
                track =>
                    track.stop()
            );


        stream =
            null;

    }


    camera.srcObject =
        null;


    startButton.disabled =
        false;


    stopButton.disabled =
        true;


    cameraPlaceholder.style.display =
        "block";


    status.textContent =
        "🟡 Camera Stopped";

}


/* =========================================================
   PROCESS CAMERA FRAME
========================================================= */

function processFrame() {

    if (!running) {

        return;

    }


    /*
        Don't send another frame while
        the previous one is being processed.
    */

    if (!frameBusy) {

        if (
            camera.readyState >= 2 &&
            camera.videoWidth > 0
        ) {

            sendFrame();

        }

    }


    requestAnimationFrame(
        processFrame
    );

}


/* =========================================================
   SEND FRAME
========================================================= */

function sendFrame() {

    frameBusy =
        true;


    const width =
        camera.videoWidth;


    const height =
        camera.videoHeight;


    /*
        Temporary canvas

        Using the camera's actual resolution.
    */

    const canvas =
        document.createElement(
            "canvas"
        );


    canvas.width =
        width;


    canvas.height =
        height;


    const ctx =
        canvas.getContext(
            "2d",
            {
                willReadFrequently:
                    true
            }
        );


    ctx.drawImage(
        camera,
        0,
        0,
        width,
        height
    );


    const imageData =
        ctx.getImageData(
            0,
            0,
            width,
            height
        );


    /*
        Copy RGBA data into an ArrayBuffer.
    */

    const rgba =
        new Uint8ClampedArray(
            imageData.data
        );


    worker.postMessage(
        {
            type:
                "frame",

            rgba:
                rgba.buffer,

            width:
                width,

            height:
                height,

            confidence:
                confidenceThreshold

        },
        [
            rgba.buffer
        ]
    );

}


/* =========================================================
   DRAW RESULT
========================================================= */

function drawResult(data) {

    const width =
        data.width;


    const height =
        data.height;


    if (
        !width ||
        !height
    ) {

        return;

    }


    output.width =
        width;


    output.height =
        height;


    /*
        Draw original frame
    */

    if (data.frame) {

        const frameBytes =
            new Uint8ClampedArray(
                data.frame
            );


        const frameImage =
            new ImageData(
                frameBytes,
                width,
                height
            );


        outputCtx.putImageData(
            frameImage,
            0,
            0
        );

    }


    /*
        Segmentation overlays
    */

    if (
        data.detections &&
        data.detections.length
    ) {

        for (
            const detection
            of data.detections
        ) {

            drawMask(
                detection,
                width,
                height
            );

        }

    }


    /*
        Draw bounding boxes + labels
    */

    if (
        data.detections
    ) {

        for (
            const detection
            of data.detections
        ) {

            drawDetection(
                detection
            );

        }

    }


    objectsText.textContent =
        data.detections
            ? data.detections.length
            : "0";


    /*
        FPS
    */

    const now =
        performance.now();


    if (lastFrameTime > 0) {

        const delta =
            now -
            lastFrameTime;


        const currentFPS =
            1000 /
            delta;


        fps =
            fps * 0.8 +
            currentFPS * 0.2;


        const fpsElement =
            document.getElementById(
                "fps"
            );


        if (fpsElement) {

            fpsElement.textContent =
                fps.toFixed(1);

        }

    }


    lastFrameTime =
        now;

}


/* =========================================================
   DRAW SEGMENTATION MASK
========================================================= */

function drawMask(
    detection,
    width,
    height
) {

    if (
        !detection.mask
    ) {

        return;

    }


    const mask =
        new Uint8Array(
            detection.mask
        );


    const maskWidth =
        detection.maskWidth;


    const maskHeight =
        detection.maskHeight;


    if (
        !maskWidth ||
        !maskHeight
    ) {

        return;

    }


    /*
        Draw mask directly onto output.

        The worker returns a mask already
        mapped to the original camera image.
    */

    const image =
        outputCtx.getImageData(
            0,
            0,
            width,
            height
        );


    const pixels =
        image.data;


    const limit =
        Math.min(
            mask.length,
            width * height
        );


    for (
        let i = 0;
        i < limit;
        i++
    ) {

        if (
            mask[i] > 0
        ) {

            const p =
                i * 4;


            /*
                Semi-transparent cyan/green
                segmentation overlay.
            */

            pixels[p] =
                Math.min(
                    255,
                    pixels[p] * 0.45 +
                    40 * 0.55
                );


            pixels[p + 1] =
                Math.min(
                    255,
                    pixels[p + 1] * 0.45 +
                    220 * 0.55
                );


            pixels[p + 2] =
                Math.min(
                    255,
                    pixels[p + 2] * 0.45 +
                    120 * 0.55
                );

        }

    }


    outputCtx.putImageData(
        image,
        0,
        0
    );

}


/* =========================================================
   DRAW BOX + LABEL
========================================================= */

function drawDetection(
    detection
) {

    const x1 =
        detection.x1;


    const y1 =
        detection.y1;


    const x2 =
        detection.x2;


    const y2 =
        detection.y2;


    const width =
        x2 - x1;


    const height =
        y2 - y1;


    outputCtx.strokeStyle =
        "#22c55e";


    outputCtx.lineWidth =
        3;


    outputCtx.strokeRect(
        x1,
        y1,
        width,
        height
    );


    const confidence =
        Number(
            detection.confidence
        ).toFixed(2);


    const label =
        `${detection.name} ${confidence}`;


    outputCtx.font =
        "bold 15px Arial";


    const textWidth =
        outputCtx.measureText(
            label
        ).width;


    const labelHeight =
        24;


    outputCtx.fillStyle =
        "#22c55e";


    outputCtx.fillRect(
        x1,
        Math.max(
            0,
            y1 - labelHeight
        ),
        textWidth + 12,
        labelHeight
    );


    outputCtx.fillStyle =
        "#020617";


    outputCtx.fillText(
        label,
        x1 + 6,
        Math.max(
            17,
            y1 - 7
        )
    );

}


/* =========================================================
   CONFIDENCE SLIDER
========================================================= */

confidence.addEventListener(
    "input",
    function() {

        confidenceThreshold =
            Number(
                confidence.value
            );


        confidenceText.textContent =
            confidenceThreshold.toFixed(
                2
            );

    }
);


/* =========================================================
   BUTTONS
========================================================= */

startButton.addEventListener(
    "click",
    startCamera
);


stopButton.addEventListener(
    "click",
    stopCamera
);


/* =========================================================
   INITIAL STATE
========================================================= */

startButton.disabled =
    true;


stopButton.disabled =
    true;


updateProgress(
    0,
    "Download the model, then select the .onnx file."
);