document.addEventListener('DOMContentLoaded', () => {
    // --- DOM Elements ---
    const uploadForm = document.getElementById('upload-form');
    const binDefinitionForm = document.getElementById('bin-definition-form');
    
    const uploadSection = document.getElementById('upload-section');
    const binDefinitionSection = document.getElementById('bin-definition-section');
    const editorControlsSection = document.getElementById('editor-controls-section');
    const editorSection = document.getElementById('editor-section');
    const welcomeMessage = document.getElementById('welcome-message');

    const uploadStatus = document.getElementById('upload-status');
    const binStatus = document.getElementById('bin-status');
    const editorStatus = document.getElementById('editor-status');

    const binMapContainer = document.getElementById('bin-map-container');
    const canvas = document.getElementById('wafer-map-canvas');
    const ctx = canvas.getContext('2d');
    
    const updateSelectionBtn = document.getElementById('update-selection-btn');
    const resetViewBtn = document.getElementById('reset-view-btn');
    const rotateBtn = document.getElementById('rotate-btn');
    const polygonSelectBtn = document.getElementById('polygon-select-btn');
    const boxSelectBtn = document.getElementById('box-select-btn');
    const insetMapBtn = document.getElementById('inset-map-btn');
    const undoBtn = document.getElementById('undo-btn');
    const redoBtn = document.getElementById('redo-btn');
    const binCountsContainer = document.getElementById('bin-counts');

    // --- State Management ---
    let waferMap = [];
    let mapRows = 0;
    let mapCols = 0;
    let selection = new Set();
    
    // --- Canvas & View State ---
    const BIN_COLORS = { 2: '#28a745', 1: '#dc3545', 0: '#6c757d', '-1': '#3e3e42' };
    const SELECTION_COLOR = 'rgba(0, 123, 255, 0.5)';
    let transform = { x: 0, y: 0, scale: 1 };
    let isPanning = false;
    let panStart = { x: 0, y: 0 };

    // --- Selection modes ---
    let selectionMode = 'box'; // 'box', 'polygon'
    let isBoxSelecting = false;
    let boxSelectionStart = { x: 0, y: 0 };
    let boxSelectionEnd = { x: 0, y: 0 };
    
    let isDrawingPolygon = false;
    let polygonPoints = [];
    let currentMousePos = { x: 0, y: 0 };

    // --- UI/UX Functions ---
    function setStatus(element, message, isError = false) {
        element.textContent = message;
        element.className = 'status';
        if (message) {
            element.classList.add(isError ? 'error' : 'success');
        }
    }

    function showSpinner(button, show = true) {
        const spinner = button.querySelector('.spinner-border');
        const btnText = button.querySelector('.btn-text');
        if (spinner) {
            spinner.style.display = show ? 'inline-block' : 'none';
            if (btnText) btnText.style.display = show ? 'none' : 'inline';
            button.disabled = show;
        }
    }

    function switchCard(from, to) {
        if(from) from.classList.remove('active');
        if(to) to.classList.add('active');
    }

    function updateHistoryButtons(canUndo, canRedo) {
        undoBtn.disabled = !canUndo;
        redoBtn.disabled = !canRedo;
    }

    function handleMapUpdate(result) {
        waferMap = result.wafer_map;
        mapRows = result.rows;
        mapCols = result.cols;
        
        selection.clear();
        updateBinCounts();
        draw();
    }

    // --- Canvas & Data Functions ---
    function worldToScreen(x, y) {
        return { x: x * transform.scale + transform.x, y: y * transform.scale + transform.y };
    }

    function screenToWorld(x, y) {
        return { x: (x - transform.x) / transform.scale, y: (y - transform.y) / transform.scale };
    }

    function getDieAtScreen(x, y) {
        const worldPos = screenToWorld(x, y);
        const col = Math.floor(worldPos.x / 10);
        const row = Math.floor(worldPos.y / 10);
        if (row >= 0 && row < mapRows && col >= 0 && col < mapCols) {
            return { row, col };
        }
        return null;
    }
    
    function resetTransform() {
        const dieSize = 10;
        const container = canvas.parentElement;
        canvas.width = container.clientWidth;
        canvas.height = container.clientHeight;

        const scaleX = canvas.width / (mapCols * dieSize);
        const scaleY = canvas.height / (mapRows * dieSize);
        transform.scale = Math.min(scaleX, scaleY) * 0.9;

        const mapRenderWidth = mapCols * dieSize * transform.scale;
        const mapRenderHeight = mapRows * dieSize * transform.scale;

        transform.x = (canvas.width - mapRenderWidth) / 2;
        transform.y = (canvas.height - mapRenderHeight) / 2;
    }

    function updateBinCounts() {
        const counts = { 2: 0, 1: 0, 0: 0, '-1': 0 };
        let total = 0;
        for (let r = 0; r < mapRows; r++) {
            for (let c = 0; c < mapCols; c++) {
                const bin = waferMap[r][c];
                if (bin in counts) counts[bin]++;
                if (bin != -1) total++;
            }
        }
        
        binCountsContainer.innerHTML = `
            <div class="count-item">Total Dies: <span>${total}</span></div>
            <div class="count-item">Good: <span style="background-color:${BIN_COLORS[2]}">${counts[2]}</span></div>
            <div class="count-item">NG: <span style="background-color:${BIN_COLORS[1]}">${counts[1]}</span></div>
            <div class="count-item">Dummy: <span style="background-color:${BIN_COLORS[0]}">${counts[0]}</span></div>
        `;
    }

    // --- Rendering Engine ---
    function draw() {
        requestAnimationFrame(() => {
            const dieSize = 10;
            ctx.clearRect(0, 0, canvas.width, canvas.height);

            ctx.save();
            ctx.translate(transform.x, transform.y);
            ctx.scale(transform.scale, transform.scale);
            
            // Draw dies
            for (let r = 0; r < mapRows; r++) {
                for (let c = 0; c < mapCols; c++) {
                    const bin = waferMap[r][c];
                    ctx.fillStyle = BIN_COLORS[bin] || '#ffffff';
                    ctx.fillRect(c * dieSize, r * dieSize, dieSize, dieSize);
                }
            }
            
            // Draw selection highlights
            ctx.fillStyle = SELECTION_COLOR;
            selection.forEach(key => {
                const [r, c] = key.split(',').map(Number);
                ctx.fillRect(c * dieSize, r * dieSize, dieSize, dieSize);
            });

            ctx.restore();
            
            // Draw selection rectangle (in screen space)
            if (isBoxSelecting) {
                ctx.fillStyle = 'rgba(0, 123, 255, 0.2)';
                ctx.strokeStyle = 'rgba(0, 123, 255, 0.8)';
                ctx.lineWidth = 1;
                const rect = {
                    x: Math.min(boxSelectionStart.x, boxSelectionEnd.x),
                    y: Math.min(boxSelectionStart.y, boxSelectionEnd.y),
                    w: Math.abs(boxSelectionStart.x - boxSelectionEnd.x),
                    h: Math.abs(boxSelectionStart.y - boxSelectionEnd.y)
                };
                ctx.fillRect(rect.x, rect.y, rect.w, rect.h);
                ctx.strokeRect(rect.x, rect.y, rect.w, rect.h);
            }

            // Draw polygon (in screen space)
            if (isDrawingPolygon && polygonPoints.length > 0) {
                ctx.beginPath();
                ctx.moveTo(polygonPoints[0].x, polygonPoints[0].y);
                for (let i = 1; i < polygonPoints.length; i++) {
                    ctx.lineTo(polygonPoints[i].x, polygonPoints[i].y);
                }
                ctx.strokeStyle = 'rgba(255, 0, 0, 0.8)';
                ctx.lineWidth = 2;
                ctx.stroke();
                
                // Draw line to current mouse position
                ctx.lineTo(currentMousePos.x, currentMousePos.y);
                ctx.strokeStyle = 'rgba(255, 0, 0, 0.5)';
                ctx.stroke();

                // Draw vertices
                ctx.fillStyle = 'red';
                polygonPoints.forEach(p => {
                    ctx.beginPath();
                    ctx.arc(p.x, p.y, 4, 0, 2 * Math.PI);
                    ctx.fill();
                });
            }
        });
    }

    // --- Event Handlers ---
    uploadForm.addEventListener('submit', async (e) => {
        e.preventDefault();
        const btn = e.submitter;
        showSpinner(btn, true);
        setStatus(uploadStatus, '');
        const formData = new FormData(uploadForm);
        try {
            const response = await fetch('/upload', { method: 'POST', body: formData });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);
            displayBinDefinition(result.unique_chars);
            switchCard(uploadSection, binDefinitionSection);
        } catch (error) {
            setStatus(uploadStatus, `Error: ${error.message}`, true);
        } finally {
            showSpinner(btn, false);
        }
    });

    function displayBinDefinition(chars) {
        binMapContainer.innerHTML = '';
        chars.forEach(char => {
            const item = document.createElement('div');
            item.className = 'bin-item';
            item.innerHTML = `
                <label>'${char}':</label>
                <select name="${char}">
                    <option value="good">Good</option>
                    <option value="ng">NG</option>
                    <option value="dummy">Dummy</option>
                    <option value="ignore">Ignore</option>
                </select>`;
            binMapContainer.appendChild(item);
        });
    }

    binDefinitionForm.addEventListener('submit', async (e) => {
        e.preventDefault();
        setStatus(binStatus, '');
        const formData = new FormData(binDefinitionForm);
        const charMap = Object.fromEntries(formData.entries());
        try {
            const response = await fetch('/generate_map', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ char_map: charMap }),
            });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);
            
            welcomeMessage.style.display = 'none';
            editorSection.style.display = 'flex';
            
            handleMapUpdate(result);
            resetTransform();
            draw();
            updateHistoryButtons(false, false); // History starts here

            switchCard(binDefinitionSection, editorControlsSection);
        } catch (error) {
            setStatus(binStatus, `Error: ${error.message}`, true);
        }
    });

    canvas.addEventListener('wheel', (e) => {
        e.preventDefault();
        if (isDrawingPolygon) return;
        const rect = canvas.getBoundingClientRect();
        const mouse = { x: e.clientX - rect.left, y: e.clientY - rect.top };
        const worldPos = screenToWorld(mouse.x, mouse.y);
        
        const zoomFactor = e.deltaY < 0 ? 1.1 : 1 / 1.1;
        transform.scale *= zoomFactor;

        const newWorldPos = screenToWorld(mouse.x, mouse.y);
        transform.x += (newWorldPos.x - worldPos.x) * transform.scale;
        transform.y += (newWorldPos.y - worldPos.y) * transform.scale;
        
        draw();
    });

    canvas.addEventListener('mousedown', (e) => {
        if (selectionMode === 'box' && e.shiftKey) {
            isBoxSelecting = true;
            boxSelectionStart = { x: e.offsetX, y: e.offsetY };
            boxSelectionEnd = boxSelectionStart;
        } else if (selectionMode === 'polygon') {
            // Polygon drawing is handled in 'click' to avoid conflict with panning
        } else {
            isPanning = true;
            panStart = { x: e.clientX, y: e.clientY };
        }
    });

    canvas.addEventListener('mousemove', (e) => {
        currentMousePos = { x: e.offsetX, y: e.offsetY };
        if (isPanning) {
            const dx = e.clientX - panStart.x;
            const dy = e.clientY - panStart.y;
            transform.x += dx;
            transform.y += dy;
            panStart = { x: e.clientX, y: e.clientY };
            draw();
        } else if (isBoxSelecting) {
            boxSelectionEnd = { x: e.offsetX, y: e.offsetY };
            draw();
        } else if (isDrawingPolygon) {
            draw();
        }
    });

    window.addEventListener('mouseup', (e) => {
        if (isPanning) isPanning = false;
        if (isBoxSelecting) {
            isBoxSelecting = false;
            selectDiesInBox();
            draw();
        }
    });
    
    canvas.addEventListener('click', (e) => {
        const wasDrag = Math.abs(e.clientX - panStart.x) > 2 || Math.abs(e.clientY - panStart.y) > 2;
        if (isPanning && wasDrag) return;

        if (selectionMode === 'polygon') {
            handlePolygonClick(e);
        } else {
            if (e.shiftKey) return; // Box selection is handled in mousedown/up
            const die = getDieAtScreen(e.offsetX, e.offsetY);
            if (die) {
                if (!e.ctrlKey) selection.clear();
                const key = `${die.row},${die.col}`;
                if (selection.has(key)) selection.delete(key);
                else selection.add(key);
                setStatus(editorStatus, `${selection.size} die(s) selected.`);
                draw();
            }
        }
    });

    canvas.addEventListener('dblclick', (e) => {
        if (selectionMode === 'polygon' && isDrawingPolygon) {
            e.preventDefault();
            closeAndSelectPolygon();
        }
    });

    // --- Selection Logic ---
    function selectDiesInBox() {
        const startWorld = screenToWorld(boxSelectionStart.x, boxSelectionStart.y);
        const endWorld = screenToWorld(boxSelectionEnd.x, boxSelectionEnd.y);
        
        const startCol = Math.floor(Math.min(startWorld.x, endWorld.x) / 10);
        const endCol = Math.floor(Math.max(startWorld.x, endWorld.x) / 10);
        const startRow = Math.floor(Math.min(startWorld.y, endWorld.y) / 10);
        const endRow = Math.floor(Math.max(startWorld.y, endWorld.y) / 10);

        if (!window.event.ctrlKey) selection.clear();
        
        for (let r = startRow; r <= endRow; r++) {
            for (let c = startCol; c <= endCol; c++) {
                if (r >= 0 && r < mapRows && c >= 0 && c < mapCols) {
                    const key = `${r},${c}`;
                    if (selection.has(key)) selection.delete(key);
                    else selection.add(key);
                }
            }
        }
        setStatus(editorStatus, `${selection.size} die(s) selected.`);
    }

    function handlePolygonClick(e) {
        const point = { x: e.offsetX, y: e.offsetY };
        if (!isDrawingPolygon) {
            isDrawingPolygon = true;
            polygonPoints = [point];
            setStatus(editorStatus, 'Polygon drawing started. Click to add points, double-click to finish.');
        } else {
            // Check if closing polygon
            const firstPoint = polygonPoints[0];
            const dist = Math.sqrt(Math.pow(point.x - firstPoint.x, 2) + Math.pow(point.y - firstPoint.y, 2));
            if (dist < 10 && polygonPoints.length > 2) {
                closeAndSelectPolygon();
            } else {
                polygonPoints.push(point);
            }
        }
        draw();
    }

    function closeAndSelectPolygon() {
        if (!window.event.ctrlKey) selection.clear();
        selectDiesInPolygon();
        isDrawingPolygon = false;
        polygonPoints = [];
        draw();
    }

    function selectDiesInPolygon() {
        const dieSize = 10;
        const polygonInWorld = polygonPoints.map(p => screenToWorld(p.x, p.y));

        for (let r = 0; r < mapRows; r++) {
            for (let c = 0; c < mapCols; c++) {
                const dieCenter = {
                    x: c * dieSize + dieSize / 2,
                    y: r * dieSize + dieSize / 2
                };
                if (pointInPolygon(dieCenter, polygonInWorld)) {
                    const key = `${r},${c}`;
                    if (selection.has(key)) selection.delete(key);
                    else selection.add(key);
                }
            }
        }
        setStatus(editorStatus, `${selection.size} die(s) selected.`);
    }

    function pointInPolygon(point, vs) {
        // ray-casting algorithm based on
        // https://wrf.ecse.rpi.edu/Research/Short_Notes/pnpoly.html
        const x = point.x, y = point.y;
        let inside = false;
        for (let i = 0, j = vs.length - 1; i < vs.length; j = i++) {
            const xi = vs[i].x, yi = vs[i].y;
            const xj = vs[j].x, yj = vs[j].y;
            const intersect = ((yi > y) !== (yj > y))
                && (x < (xj - xi) * (y - yi) / (yj - yi) + xi);
            if (intersect) inside = !inside;
        }
        return inside;
    }

    // --- Action Buttons ---
    function setSelectionMode(mode) {
        selectionMode = mode;
        if (mode === 'polygon') {
            polygonSelectBtn.classList.add('active');
            boxSelectBtn.classList.remove('active');
            setStatus(editorStatus, 'Polygon select mode. Click to draw, dbl-click to finish.');
            isDrawingPolygon = false;
            polygonPoints = [];
        } else { // box mode
            boxSelectBtn.classList.add('active');
            polygonSelectBtn.classList.remove('active');
            setStatus(editorStatus, 'Box select mode. Hold Shift and drag to select.');
            isDrawingPolygon = false; // Cancel any ongoing polygon drawing
        }
        draw();
    }

    boxSelectBtn.addEventListener('click', () => setSelectionMode('box'));
    polygonSelectBtn.addEventListener('click', () => setSelectionMode('polygon'));

    updateSelectionBtn.addEventListener('click', async () => {
        if (selection.size === 0) {
            setStatus(editorStatus, 'No dies selected to update.', true);
            return;
        }
        const binType = document.getElementById('bin-type-selector').value;
        const diesToUpdate = Array.from(selection).map(key => key.split(',').map(Number));
        
        setStatus(editorStatus, `Updating ${selection.size} die(s)...`);

        try {
            const response = await fetch('/update_map', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ dies: diesToUpdate, bin_type: binType }),
            });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);

            handleMapUpdate(result);
            setStatus(editorStatus, 'Update successful.');
            updateHistoryButtons(true, false);
        } catch (error) {
            setStatus(editorStatus, `Error: ${error.message}`, true);
        }
    });
    
    resetViewBtn.addEventListener('click', async () => {
        setStatus(editorStatus, 'Resetting map to original state...');
        try {
            const response = await fetch('/reset_map', { method: 'POST' });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);
            
            handleMapUpdate(result);
            resetTransform(); // Also reset zoom and pan
            draw();
            setStatus(editorStatus, 'Map has been reset.');
            updateHistoryButtons(result.can_undo, result.can_redo);

        } catch (error) {
            setStatus(editorStatus, `Error: ${error.message}`, true);
        }
    });

    rotateBtn.addEventListener('click', async () => {
        setStatus(editorStatus, 'Rotating map...');
        try {
            const response = await fetch('/rotate_map', { method: 'POST' });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);

            handleMapUpdate(result);
            setStatus(editorStatus, 'Map rotated successfully.');
            updateHistoryButtons(true, false);
        } catch (error) {
            setStatus(editorStatus, `Error: ${error.message}`, true);
        }
    });

    insetMapBtn.addEventListener('click', async () => {
        const layers = document.getElementById('inset-layers').value;
        const fromBin = document.getElementById('inset-from-bin').value;
        const toBin = document.getElementById('inset-to-bin').value;

        if (parseInt(layers) < 1) {
            setStatus(editorStatus, 'Inset layers must be at least 1.', true);
            return;
        }
        if (fromBin === toBin) {
            setStatus(editorStatus, "'From' and 'To' bin types cannot be the same.", true);
            return;
        }

        setStatus(editorStatus, `Insetting ${layers} layer(s)...`);

        try {
            const response = await fetch('/inset_map', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    layers: layers,
                    from_bin: fromBin,
                    to_bin: toBin
                }),
            });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);

            handleMapUpdate(result);
            setStatus(editorStatus, 'Inset operation successful.');
            updateHistoryButtons(true, false);
        } catch (error) {
            setStatus(editorStatus, `Error: ${error.message}`, true);
        }
    });

    async function handleHistoryAction(action) {
        setStatus(editorStatus, `Performing ${action}...`);
        try {
            const response = await fetch(`/${action}`, { method: 'POST' });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);
            
            handleMapUpdate(result);
            setStatus(editorStatus, `${action} successful.`);
            updateHistoryButtons(result.can_undo, result.can_redo);

        } catch (error) {
            setStatus(editorStatus, `Error: ${error.message}`, true);
        }
    }

    undoBtn.addEventListener('click', () => handleHistoryAction('undo'));
    redoBtn.addEventListener('click', () => handleHistoryAction('redo'));
    
    new ResizeObserver(() => {
        if (mapRows > 0) { // Only resize if there is a map
            resetTransform();
            draw();
        }
    }).observe(canvas.parentElement);
});