document.addEventListener('DOMContentLoaded', () => {
    // --- DOM Elements ---
    const uploadForm = document.getElementById('upload-form');
    const binDefinitionForm = document.getElementById('bin-definition-form');
    const fileInput = document.getElementById('file-input');

    const uploadSection = document.getElementById('upload-section');
    const binDefinitionSection = document.getElementById('bin-definition-section');
    const comparisonControlsSection = document.getElementById('comparison-controls-section');
    const viewerSection = document.getElementById('viewer-section');
    const comparisonResultSection = document.getElementById('comparison-result-section');
    const welcomeMessage = document.getElementById('welcome-message');

    const uploadStatus = document.getElementById('upload-status');
    const binStatus = document.getElementById('bin-status');
    const compareStatus = document.getElementById('compare-status');

    const binMapContainer = document.getElementById('bin-map-container');
    const referenceSetterContainer = document.getElementById('reference-setter-container');
    
    const mapSelector = document.getElementById('map-selector');
    const rotateMapBtn = document.getElementById('rotate-map-btn');
    const compareBtn = document.getElementById('compare-btn');
    const backToViewerBtn = document.getElementById('back-to-viewer-btn');

    const canvas = document.getElementById('wafer-map-canvas');
    const ctx = canvas.getContext('2d');
    const resultCanvas = document.getElementById('result-map-canvas');
    const resultCtx = resultCanvas.getContext('2d');
    const resultLegendContainer = document.getElementById('result-legend');
    const tooltip = document.getElementById('map-tooltip');

    // --- State Management ---
    let sessionMaps = {};
    let activeMapFid = null;
    let comparisonData = {}; // To store result_map, detailed_grid, etc.

    // --- Canvas & View State ---
    const BIN_CODE_MAP = { 2: 'Good', 1: 'NG', 0: 'Dummy', '-1': 'Ignore', '-10': 'N/A' };
    const BIN_COLORS = { 2: '#28a745', 1: '#dc3545', 0: '#6c757d', '-1': '#3e3e42' };
    const REFERENCE_COLOR = 'rgba(255, 193, 7, 1)';
    let transform = { x: 0, y: 0, scale: 1 };
    let isPanning = false;
    let panStart = { x: 0, y: 0 };
    
    let resultTransform = { x: 0, y: 0, scale: 1 };
    let isResultPanning = false;
    let resultPanStart = { x: 0, y: 0 };
    let highlightedDie = null;

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

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

    function updateReferenceUI() {
        referenceSetterContainer.innerHTML = '';
        Object.keys(sessionMaps).forEach(fid => {
            const map = sessionMaps[fid];
            const hasRef = map.reference !== null;
            const refText = hasRef ? `(${map.reference.join(', ')})` : 'Not Set';
            
            const item = document.createElement('div');
            item.className = `reference-item ${hasRef ? 'has-ref' : ''}`;
            item.id = `ref-item-${fid}`;
            item.innerHTML = `
                <span class="filename" title="${map.filename}">${map.filename}</span>
                <div class="ref-details">
                    <span class="ref-status ${hasRef ? 'set' : ''}">${refText}</span>
                    <button class="clear-ref-btn" data-fid="${fid}" title="Clear reference">×</button>
                </div>
            `;
            referenceSetterContainer.appendChild(item);
        });
        checkAllReferencesSet();
    }
    
    function checkAllReferencesSet() {
        const allSet = Object.values(sessionMaps).every(m => m.reference !== null);
        compareBtn.disabled = !allSet;
        if (allSet) {
            setStatus(compareStatus, 'All references set. Ready to compare.');
        } else {
            setStatus(compareStatus, 'Please set a reference point for each map.');
        }
    }

    // --- Canvas & Drawing ---
    function worldToScreen(x, y, t) { return { x: x * t.scale + t.x, y: y * t.scale + t.y }; }
    function screenToWorld(x, y, t) { return { x: (x - t.x) / t.scale, y: (y - t.y) / t.scale }; }

    function resetTransform(map, t, targetCanvas) {
        const dieSize = 10;
        const container = targetCanvas.parentElement;
        targetCanvas.width = container.clientWidth;
        targetCanvas.height = container.clientHeight;

        const scaleX = targetCanvas.width / (map.cols * dieSize);
        const scaleY = targetCanvas.height / (map.rows * dieSize);
        t.scale = Math.min(scaleX, scaleY) * 0.9;

        const mapRenderWidth = map.cols * dieSize * t.scale;
        const mapRenderHeight = map.rows * dieSize * t.scale;

        t.x = (targetCanvas.width - mapRenderWidth) / 2;
        t.y = (targetCanvas.height - mapRenderHeight) / 2;
    }

    function drawActiveMap() {
        if (!activeMapFid || !sessionMaps[activeMapFid] || !sessionMaps[activeMapFid].wafer_map) return;
        const map = sessionMaps[activeMapFid];
        
        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);

            for (let r = 0; r < map.rows; r++) {
                for (let c = 0; c < map.cols; c++) {
                    const bin = map.wafer_map[r][c];
                    ctx.fillStyle = BIN_COLORS[bin] || '#ffffff';
                    ctx.fillRect(c * dieSize, r * dieSize, dieSize, dieSize);
                }
            }

            if (map.reference) {
                const [r, c] = map.reference;
                ctx.strokeStyle = REFERENCE_COLOR;
                ctx.lineWidth = 2 / transform.scale;
                ctx.strokeRect(c * dieSize - 1, r * dieSize - 1, dieSize + 2, dieSize + 2);
            }
            ctx.restore();
        });
    }
    
    function drawResultMap() {
        const { rows, cols, result_map, legend } = comparisonData;
        
        requestAnimationFrame(() => {
            const dieSize = 10;
            resultCtx.clearRect(0, 0, resultCanvas.width, resultCanvas.height);
            resultCtx.save();
            resultCtx.translate(resultTransform.x, resultTransform.y);
            resultCtx.scale(resultTransform.scale, resultTransform.scale);

            for (let r = 0; r < rows; r++) {
                for (let c = 0; c < cols; c++) {
                    const code = result_map[r][c];
                    resultCtx.fillStyle = legend[code] ? legend[code].color : '#ffffff';
                    resultCtx.fillRect(c * dieSize, r * dieSize, dieSize, dieSize);
                }
            }
            
            // Draw highlight on hovered die
            if (highlightedDie) {
                const { r, c } = highlightedDie;
                resultCtx.strokeStyle = REFERENCE_COLOR;
                resultCtx.lineWidth = 2 / resultTransform.scale;
                resultCtx.strokeRect(c * dieSize, r * dieSize, dieSize, dieSize);
            }

            resultCtx.restore();
        });
    }

    // --- Event Handlers ---
    uploadForm.addEventListener('submit', async (e) => {
        e.preventDefault();
        setStatus(uploadStatus, '');
        if (fileInput.files.length < 2) {
            setStatus(uploadStatus, 'Please select at least two files for comparison.', true);
            return;
        }
        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);
            
            binMapContainer.innerHTML = '';
            result.unique_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);
            });
            
            mapSelector.innerHTML = '';
            sessionMaps = {}; // Clear previous session
            Object.keys(result.maps).forEach(fid => {
                sessionMaps[fid] = { filename: result.maps[fid], wafer_map: null, reference: null };
                const option = document.createElement('option');
                option.value = fid;
                option.textContent = result.maps[fid];
                mapSelector.appendChild(option);
            });
            activeMapFid = mapSelector.value;

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

    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);

            Object.keys(result.maps_data).forEach(fid => {
                sessionMaps[fid] = { ...sessionMaps[fid], ...result.maps_data[fid] };
            });
            
            welcomeMessage.style.display = 'none';
            viewerSection.style.display = 'flex';
            
            resetTransform(sessionMaps[activeMapFid], transform, canvas);
            drawActiveMap();
            updateReferenceUI();
            switchCard(binDefinitionSection, comparisonControlsSection);
        } catch (error) {
            setStatus(binStatus, `Error: ${error.message}`, true);
        }
    });

    mapSelector.addEventListener('change', (e) => {
        activeMapFid = e.target.value;
        resetTransform(sessionMaps[activeMapFid], transform, canvas);
        drawActiveMap();
    });

    rotateMapBtn.addEventListener('click', async () => {
        if (!activeMapFid) return;
        try {
            const response = await fetch('/compare/rotate', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ fid: activeMapFid }),
            });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);

            sessionMaps[activeMapFid].wafer_map = result.wafer_map;
            sessionMaps[activeMapFid].rows = result.rows;
            sessionMaps[activeMapFid].cols = result.cols;
            sessionMaps[activeMapFid].reference = null;

            resetTransform(sessionMaps[activeMapFid], transform, canvas);
            drawActiveMap();
            updateReferenceUI();
        } catch (error) {
            setStatus(compareStatus, `Error rotating map: ${error.message}`, true);
        }
    });

    referenceSetterContainer.addEventListener('click', async (e) => {
        if (e.target.classList.contains('clear-ref-btn')) {
            const fid = e.target.dataset.fid;
            try {
                const response = await fetch('/compare/clear_reference', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ fid: fid }),
                });
                const result = await response.json();
                if (!response.ok) throw new Error(result.error);

                sessionMaps[fid].reference = null;
                updateReferenceUI();
                if (fid === activeMapFid) drawActiveMap();

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

    canvas.addEventListener('click', async (e) => {
        const wasDrag = Math.abs(e.clientX - panStart.x) > 2 || Math.abs(e.clientY - panStart.y) > 2;
        if ((isPanning && wasDrag) || !activeMapFid) return;

        const worldPos = screenToWorld(e.offsetX, e.offsetY, transform);
        const col = Math.floor(worldPos.x / 10);
        const row = Math.floor(worldPos.y / 10);
        
        const map = sessionMaps[activeMapFid];
        if (row >= 0 && row < map.rows && col >= 0 && col < map.cols) {
            const reference = [row, col];
            try {
                const response = await fetch('/set_reference', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ fid: activeMapFid, reference: reference }),
                });
                const result = await response.json();
                if (!response.ok) throw new Error(result.error);

                sessionMaps[activeMapFid].reference = result.reference;
                updateReferenceUI();
                drawActiveMap();

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

    compareBtn.addEventListener('click', async () => {
        setStatus(compareStatus, 'Running comparison...');
        try {
            const response = await fetch('/run_comparison', { method: 'POST' });
            const result = await response.json();
            if (!response.ok) throw new Error(result.error);
            
            comparisonData = result;
            viewerSection.style.display = 'none';
            comparisonResultSection.style.display = 'flex';
            
            resetTransform(comparisonData, resultTransform, resultCanvas);
            drawResultMap();
            
            resultLegendContainer.innerHTML = '<h3>Legend</h3>';
            for (const code in result.legend) {
                const item = result.legend[code];
                const legendItem = document.createElement('div');
                legendItem.className = 'count-item';
                legendItem.innerHTML = `${item.label}: <span style="background-color:${item.color}"></span>`;
                resultLegendContainer.appendChild(legendItem);
            }

        } catch (error) {
            setStatus(compareStatus, `Error: ${error.message}`, true);
        }
    });
    
    backToViewerBtn.addEventListener('click', () => {
        comparisonResultSection.style.display = 'none';
        viewerSection.style.display = 'flex';
    });

    // --- Canvas Pan & Zoom ---
    function setupCanvasEvents(targetCanvas, t, panFlagSetter, drawFn) {
        targetCanvas.addEventListener('wheel', (e) => { e.preventDefault(); const rect = targetCanvas.getBoundingClientRect(); const mouse = { x: e.clientX - rect.left, y: e.clientY - rect.top }; const worldPos = screenToWorld(mouse.x, mouse.y, t); const zoomFactor = e.deltaY < 0 ? 1.1 : 1 / 1.1; t.scale *= zoomFactor; const newWorldPos = screenToWorld(mouse.x, mouse.y, t); t.x += (newWorldPos.x - worldPos.x) * t.scale; t.y += (newWorldPos.y - worldPos.y) * t.scale; drawFn(); });
        targetCanvas.addEventListener('mousedown', (e) => { panFlagSetter(true, e); });
        targetCanvas.addEventListener('mousemove', (e) => { panFlagSetter(false, e, true, drawFn); });
    }
    
    setupCanvasEvents(canvas, transform, (val, e) => { isPanning = val; if(val) panStart = { x: e.clientX, y: e.clientY }; }, drawActiveMap);
    setupCanvasEvents(resultCanvas, resultTransform, (val, e) => { isResultPanning = val; if(val) resultPanStart = { x: e.clientX, y: e.clientY }; }, drawResultMap);
    
    window.addEventListener('mouseup', () => { isPanning = false; isResultPanning = false; });
    window.addEventListener('mousemove', (e) => {
        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 }; drawActiveMap(); }
        if (isResultPanning) { const dx = e.clientX - resultPanStart.x; const dy = e.clientY - resultPanStart.y; resultTransform.x += dx; resultTransform.y += dy; resultPanStart = { x: e.clientX, y: e.clientY }; drawResultMap(); }
    });

    // --- Tooltip Logic ---
    resultCanvas.addEventListener('mousemove', (e) => {
        if (!comparisonData.detailed_grid || isResultPanning) return;
        
        const rect = resultCanvas.getBoundingClientRect();
        const mouse = { x: e.clientX - rect.left, y: e.clientY - rect.top };
        const worldPos = screenToWorld(mouse.x, mouse.y, resultTransform);
        const c = Math.floor(worldPos.x / 10);
        const r = Math.floor(worldPos.y / 10);

        if (r >= 0 && r < comparisonData.rows && c >= 0 && c < comparisonData.cols) {
            highlightedDie = { r, c };
            const dieData = comparisonData.detailed_grid[r][c];
            let tooltipContent = `Coord: (${r}, ${c})\n-------------------\n`;
            tooltipContent += comparisonData.map_names.map((name, i) => {
                const binCode = dieData[i];
                return `${name}: ${BIN_CODE_MAP[binCode] || 'Unknown'}`;
            }).join('\n');
            
            tooltip.style.display = 'block';
            tooltip.style.opacity = '1';
            tooltip.textContent = tooltipContent;
            
            // Position tooltip
            const tooltipRect = tooltip.getBoundingClientRect();
            let left = e.clientX + 15;
            let top = e.clientY + 15;
            if (left + tooltipRect.width > window.innerWidth) {
                left = e.clientX - tooltipRect.width - 15;
            }
            if (top + tooltipRect.height > window.innerHeight) {
                top = e.clientY - tooltipRect.height - 15;
            }
            tooltip.style.left = `${left}px`;
            tooltip.style.top = `${top}px`;

        } else {
            highlightedDie = null;
            tooltip.style.display = 'none';
            tooltip.style.opacity = '0';
        }
        drawResultMap();
    });
    resultCanvas.addEventListener('mouseleave', () => {
        highlightedDie = null;
        tooltip.style.display = 'none';
        tooltip.style.opacity = '0';
        drawResultMap();
    });

    new ResizeObserver(() => { if (activeMapFid) { resetTransform(sessionMaps[activeMapFid], transform, canvas); drawActiveMap(); } }).observe(canvas.parentElement);
    new ResizeObserver(() => { if (comparisonData.rows) { resetTransform(comparisonData, resultTransform, resultCanvas); drawResultMap(); } }).observe(resultCanvas.parentElement);
});