/**
 * ArtiAgent Web UI - Unified Frontend Logic
 * Supports both DefectFill & DefectDiffu
 * ==========================================
 */

// =============================================================================
// Global State
// =============================================================================

const state = {
    currentModel: 'defectfill',
    activeTab: 'training',
    uploadedImages: [],
    currentMaskImage: null,
    drawnRect: null,
    drawnRectNatural: null,
    isDrawing: false,
    trainingPollInterval: null,
    checkpoints: {},
    generationResults: [],
    savedMasks: new Set(),
    ddConfig: {},
};

// =============================================================================
// Initialization
// =============================================================================

document.addEventListener('DOMContentLoaded', () => {
    initTabs();
    initModelSelector();
    initTrainingWizard();
    initGenerationForm();
    initDefectDiffu();
    initCanvas();
    loadCheckpoints();
    loadSystemInfo();
    loadDDConfig();
    setupDomainExamples();
});

// =============================================================================
// Navigation & Tabs
// =============================================================================

function initTabs() {
    const tabBtns = document.querySelectorAll('.tab-btn');
    const tabContents = document.querySelectorAll('.tab-content');

    tabBtns.forEach(btn => {
        btn.addEventListener('click', () => {
            const target = btn.dataset.tab;
            state.activeTab = target;

            tabBtns.forEach(b => b.classList.remove('active'));
            tabContents.forEach(c => c.classList.remove('active'));

            btn.classList.add('active');
            document.getElementById(`tab-${target}`).classList.add('active');
        });
    });
}

function initModelSelector() {
    const modelBtns = document.querySelectorAll('.model-btn');
    modelBtns.forEach(btn => {
        btn.addEventListener('click', () => {
            const model = btn.dataset.model;
            state.currentModel = model;

            modelBtns.forEach(b => b.classList.remove('active'));
            btn.classList.add('active');

            const defectfillTabs = document.getElementById('defectfill-tabs');
            const defectdiffuTabs = document.getElementById('defectdiffu-tabs');

            if (model === 'defectfill') {
                defectfillTabs.classList.remove('hidden');
                defectdiffuTabs.classList.add('hidden');
            } else {
                defectfillTabs.classList.add('hidden');
                defectdiffuTabs.classList.remove('hidden');
            }
        });
    });
}

// =============================================================================
// Training Wizard (DefectFill)
// =============================================================================

function initTrainingWizard() {
    const btnStep1 = document.getElementById('btn-step1');
    btnStep1.addEventListener('click', () => {
        const objClass = document.getElementById('train-object-class').value.trim();
        const defectType = document.getElementById('train-defect-type').value.trim();

        if (!objClass || !defectType) {
            showAlert('step1-alert', 'Please enter both product name and defect type.', 'danger');
            return;
        }

        document.getElementById('step2').classList.remove('hidden');
        document.getElementById('step2').scrollIntoView({ behavior: 'smooth' });
        checkGoodImages(objClass);
    });

    const goodUploadZone = document.getElementById('good-upload-zone');
    const goodFileInput = document.getElementById('good-images');

    goodUploadZone.addEventListener('click', () => goodFileInput.click());
    goodUploadZone.addEventListener('dragover', (e) => {
        e.preventDefault();
        goodUploadZone.classList.add('dragover');
    });
    goodUploadZone.addEventListener('dragleave', () => goodUploadZone.classList.remove('dragover'));
    goodUploadZone.addEventListener('drop', (e) => {
        e.preventDefault();
        goodUploadZone.classList.remove('dragover');
        handleGoodFileSelect(e.dataTransfer.files);
    });

    goodFileInput.addEventListener('change', (e) => handleGoodFileSelect(e.target.files));
    document.getElementById('btn-upload-good').addEventListener('click', () => {
        const files = goodFileInput.files;
        if (files.length === 0) {
            showAlert('step2-alert', 'Please select files first.', 'warning');
            return;
        }
        handleGoodFileSelect(files);
    });
    document.getElementById('btn-use-existing-good').addEventListener('click', () => {
        document.getElementById('step3').classList.remove('hidden');
        document.getElementById('step3').scrollIntoView({ behavior: 'smooth' });
    });
    document.getElementById('btn-proceed-good').addEventListener('click', () => {
        document.getElementById('step3').classList.remove('hidden');
        document.getElementById('step3').scrollIntoView({ behavior: 'smooth' });
    });

    const defectUploadZone = document.getElementById('defect-upload-zone');
    const defectFileInput = document.getElementById('defect-images');

    defectUploadZone.addEventListener('click', () => defectFileInput.click());
    defectUploadZone.addEventListener('dragover', (e) => {
        e.preventDefault();
        defectUploadZone.classList.add('dragover');
    });
    defectUploadZone.addEventListener('dragleave', () => defectUploadZone.classList.remove('dragover'));
    defectUploadZone.addEventListener('drop', (e) => {
        e.preventDefault();
        defectUploadZone.classList.remove('dragover');
        handleFileSelect(e.dataTransfer.files);
    });

    defectFileInput.addEventListener('change', (e) => handleFileSelect(e.target.files));
    document.getElementById('btn-step3-upload').addEventListener('click', () => {
        const files = defectFileInput.files;
        if (files.length === 0) {
            showAlert('step3-alert', 'Please select files first.', 'warning');
            return;
        }
        handleFileSelect(files);
    });

    document.getElementById('btn-align-yes').addEventListener('click', () => {
        document.getElementById('align-reference-section').classList.remove('hidden');
    });
    document.getElementById('btn-align-no').addEventListener('click', skipAlignment);

    const refUploadZone = document.getElementById('reference-upload-zone');
    const refFileInput = document.getElementById('reference-image');
    refUploadZone.addEventListener('click', () => refFileInput.click());
    refFileInput.addEventListener('change', (e) => {
        if (e.target.files && e.target.files.length > 0) {
            renderReferenceImage(e.target.files[0]);
        }
    });
    document.getElementById('btn-align-proceed').addEventListener('click', runAlignment);

    document.getElementById('btn-save-mask').addEventListener('click', saveMask);
    document.getElementById('btn-skip-mask').addEventListener('click', skipMask);
    document.getElementById('btn-prev-image').addEventListener('click', prevMaskImage);
    document.getElementById('btn-next-image').addEventListener('click', nextMaskImage);
    document.getElementById('btn-clear-rect').addEventListener('click', clearRect);
    document.getElementById('btn-done-masks').addEventListener('click', finishMaskStep);

    document.getElementById('btn-save-domain').addEventListener('click', saveDomainInstructions);
    document.getElementById('btn-load-domain-example').addEventListener('click', loadDomainExample);

    document.getElementById('btn-start-training').addEventListener('click', startTraining);
    document.getElementById('btn-stop-training').addEventListener('click', stopTraining);
}

function handleFileSelect(files) {
    const objClass = document.getElementById('train-object-class').value.trim();
    const defectType = document.getElementById('train-defect-type').value.trim();

    if (!objClass || !defectType) {
        showAlert('step3-alert', 'Please complete Step 1 first.', 'warning');
        return;
    }

    if (files.length < 1) {
        showAlert('step3-alert', 'At least 1 images required for few-shot training.', 'warning');
        return;
    }

    const formData = new FormData();
    formData.append('object_class', objClass);
    formData.append('defect_type', defectType);

    for (const file of files) {
        if (file.type.startsWith('image/')) {
            formData.append('images', file);
        }
    }

    const btn = document.getElementById('btn-step3-upload');
    btn.disabled = true;
    btn.innerHTML = '<span class="spinner spinner-dark"></span> Uploading...';

    fetch('/api/training/upload-images', {
        method: 'POST',
        body: formData
    })
    .then(r => r.json())
    .then(data => {
        btn.disabled = false;
        btn.innerHTML = '📤 Upload Images';

        if (data.success) {
            state.uploadedImages = data.saved;
            renderUploadedImages();
            showAlert('step3-alert', `Uploaded ${data.count} images successfully.`, 'success');
            document.getElementById('step3-5').classList.remove('hidden');
            document.getElementById('step3-5').scrollIntoView({ behavior: 'smooth' });
        } else {
            showAlert('step3-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        btn.innerHTML = '📤 Upload Images';
        showAlert('step3-alert', 'Upload failed: ' + err.message, 'danger');
    });
}

function checkProductExists(objClass, defectType) {
    fetch(`/api/check-product-exists?object_class=${encodeURIComponent(objClass)}`)
    .then(r => r.json())
    .then(data => {
        if (data.success && data.exists) {
            showDuplicateProductModal(objClass, data.existing_defect_types, data.sample_image_url, defectType);
        } else {
            proceedToStep2();
        }
    })
    .catch(err => {
        console.error('Error checking product existence:', err);
        proceedToStep2();
    });
}

function showDuplicateProductModal(productName, existingDefectTypes, sampleImageUrl, defectType) {
    const modal = document.getElementById('duplicate-product-modal');
    document.getElementById('dup-modal-product-name').textContent = productName;
    document.getElementById('dup-modal-defect-types').textContent =
        existingDefectTypes.length > 0 ? existingDefectTypes.join(', ') : 'None';

    const imgContainer = document.getElementById('dup-modal-image-container');
    const img = document.getElementById('dup-modal-sample-image');
    if (sampleImageUrl) {
        img.src = sampleImageUrl;
        imgContainer.classList.remove('hidden');
    } else {
        imgContainer.classList.add('hidden');
    }

    modal.classList.remove('hidden');

    document.getElementById('btn-add-defect-type').onclick = () => {
        modal.classList.add('hidden');
        proceedToStep2();
    };

    document.getElementById('btn-new-product').onclick = () => {
        modal.classList.add('hidden');
        promptForNewProductName(productName);
    };

    document.getElementById('btn-close-dup-modal').onclick = () => {
        modal.classList.add('hidden');
    };

    modal.onclick = (e) => {
        if (e.target === modal) modal.classList.add('hidden');
    };
}

function promptForNewProductName(originalName) {
    const newName = prompt(
        `Product "${originalName}" already exists.\n\nPlease enter a new, unique product name:`
    );
    if (!newName || !newName.trim()) {
        showAlert('step1-alert', 'A new product name is required to proceed.', 'warning');
        return;
    }
    const trimmed = newName.trim();

    fetch(`/api/check-product-exists?object_class=${encodeURIComponent(trimmed)}`)
    .then(r => r.json())
    .then(data => {
        if (data.success && data.exists) {
            showAlert('step1-alert', `Product name "${trimmed}" also already exists. Please choose another name.`, 'danger');
            setTimeout(() => promptForNewProductName(originalName), 300);
        } else {
            document.getElementById('train-object-class').value = trimmed;
            showAlert('step1-alert', `Product name updated to "${trimmed}". Proceeding...`, 'success');
            proceedToStep2();
        }
    })
    .catch(err => {
        showAlert('step1-alert', 'Error checking new name: ' + err.message, 'danger');
    });
}

function checkGoodImages(objectClass) {
    fetch(`/api/training/good-images/count?object_class=${encodeURIComponent(objectClass)}`)
    .then(r => r.json())
    .then(data => {
        const statusDiv = document.getElementById('good-images-status');
        const useExistingBtn = document.getElementById('btn-use-existing-good');
        const proceedBtn = document.getElementById('btn-proceed-good');

        if (data.success && data.count > 0) {
            statusDiv.className = 'alert alert-info';
            statusDiv.innerHTML = `<strong>${data.count}</strong> good image(s) already exist for "<strong>${escapeHtml(objectClass)}</strong>".<br>You can upload more or proceed with existing ones.`;
            statusDiv.classList.remove('hidden');
            useExistingBtn.classList.remove('hidden');
            proceedBtn.classList.remove('hidden');
        } else {
            statusDiv.className = 'alert alert-warning';
            statusDiv.innerHTML = `No good images found for "<strong>${escapeHtml(objectClass)}</strong>".<br>Please upload at least 1 clean (non-defective) image.`;
            statusDiv.classList.remove('hidden');
            useExistingBtn.classList.add('hidden');
            proceedBtn.classList.add('hidden');
        }
    })
    .catch(err => {
        showAlert('step2-alert', 'Failed to check good images: ' + err.message, 'danger');
    });
}

function handleGoodFileSelect(files) {
    const objClass = document.getElementById('train-object-class').value.trim();

    if (!objClass) {
        showAlert('step2-alert', 'Please complete Step 1 first.', 'warning');
        return;
    }

    const formData = new FormData();
    formData.append('object_class', objClass);

    for (const file of files) {
        if (file.type.startsWith('image/')) {
            formData.append('images', file);
        }
    }

    const btn = document.getElementById('btn-upload-good');
    btn.disabled = true;
    btn.innerHTML = '<span class="spinner spinner-dark"></span> Uploading...';

    fetch('/api/training/upload-good-images', {
        method: 'POST',
        body: formData
    })
    .then(r => r.json())
    .then(data => {
        btn.disabled = false;
        btn.innerHTML = '📤 Upload Good Images';

        if (data.success) {
            renderGoodImages(data.saved);
            showAlert('step2-alert', `Uploaded ${data.count} good images. Total in folder: ${data.total_in_folder}`, 'success');
            document.getElementById('btn-proceed-good').classList.remove('hidden');
            document.getElementById('btn-use-existing-good').classList.remove('hidden');
        } else {
            showAlert('step2-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        btn.innerHTML = '📤 Upload Good Images';
        showAlert('step2-alert', 'Upload failed: ' + err.message, 'danger');
    });
}

function renderGoodImages(images) {
    const gallery = document.getElementById('good-images-gallery');
    images.forEach(img => {
        const div = document.createElement('div');
        div.className = 'image-item';
        div.innerHTML = `<img src="${img.url}" alt="${img.filename}">`;
        gallery.appendChild(div);
    });
}

function proceedToStep2() {
    document.getElementById('step2').classList.remove('hidden');
    document.getElementById('step2').scrollIntoView({ behavior: 'smooth' });
}

function renderUploadedImages() {
    const gallery = document.getElementById('defect-upload-gallery');
    gallery.innerHTML = '';

    state.uploadedImages.forEach((img, idx) => {
        const div = document.createElement('div');
        div.className = 'image-item';
        div.innerHTML = `
            <img src="${img.url}" alt="${img.filename}">
            <span class="badge-mask hidden" id="mask-badge-${idx}">MASK</span>
            <div class="overlay">
                <button class="btn btn-sm btn-primary" onclick="initMaskCanvas(${idx})">Draw Mask</button>
            </div>
        `;
        gallery.appendChild(div);
    });

    updateMaskCounter();
}

function renderReferenceImage(file) {
    const gallery = document.getElementById('reference-gallery');
    gallery.innerHTML = '';
    const div = document.createElement('div');
    div.className = 'image-item';
    const url = URL.createObjectURL(file);
    div.innerHTML = `<img src="${url}" alt="Reference">`;
    gallery.appendChild(div);
}

function skipAlignment() {
    const objClass = document.getElementById('train-object-class').value.trim();
    const defectType = document.getElementById('train-defect-type').value.trim();
    const btn = document.getElementById('btn-align-no');
    btn.disabled = true;

    fetch('/api/training/skip-align', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ object_class: objClass })
    })
    .then(r => r.json())
    .then(data => {
        btn.disabled = false;
        if (data.success) {
            if (data.images) {
                const defectImages = data.images.filter(img => 
                    img.path.includes(`/defective/${encodeURIComponent(defectType)}`)
                );

                if (defectImages.length > 0) {
                    state.uploadedImages = defectImages;
                    renderUploadedImages();
                }
            }
            showAlert('step3-5-alert', `Moved ${data.moved_count} images to final location.`, 'success');
            proceedToStep4();
        } else {
            showAlert('step3-5-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        showAlert('step3-5-alert', 'Error: ' + err.message, 'danger');
    });
}

function runAlignment() {
    const objClass = document.getElementById('train-object-class').value.trim();
    const defectType = document.getElementById('train-defect-type').value.trim();
    const refInput = document.getElementById('reference-image');

    if (!refInput.files || refInput.files.length === 0) {
        showAlert('step3-5-alert', 'Please upload a reference image first.', 'warning');
        return;
    }

    const formData = new FormData();
    formData.append('object_class', objClass);
    formData.append('reference', refInput.files[0]);

    const btn = document.getElementById('btn-align-proceed');
    const spinner = document.getElementById('align-spinner');
    btn.disabled = true;
    spinner.classList.remove('hidden');

    fetch('/api/training/align', {
        method: 'POST',
        body: formData
    })
    .then(r => r.json())
    .then(data => {
        btn.disabled = false;
        spinner.classList.add('hidden');

        if (data.success) {
            if (data.images) {
                const defectImages = data.images.filter(img => 
                    img.path.includes(`/defective/${encodeURIComponent(defectType)}`)
                );

                if (defectImages.length > 0) {
                    state.uploadedImages = defectImages;
                    renderUploadedImages();
                }
            }

            let msg = `Aligned ${data.aligned_count} images successfully.`;
            if (data.error_count > 0) {
                msg += ` ${data.error_count} failed and were copied as-is.`;
            }
            showAlert('step3-5-alert', msg, 'success');
            proceedToStep4();
        } else {
            showAlert('step3-5-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        spinner.classList.add('hidden');
        showAlert('step3-5-alert', 'Alignment failed: ' + err.message, 'danger');
    });
}

function proceedToStep4() {
    document.getElementById('step4').classList.remove('hidden');
    document.getElementById('step4').scrollIntoView({ behavior: 'smooth' });
    initMaskCanvas(0);
}

// =============================================================================
// Canvas Mask Drawing
// =============================================================================

function initCanvas() {
    // Canvas is initialized per-image
}

function initMaskCanvas(index) {
    if (index < 0 || index >= state.uploadedImages.length) return;

    state.currentMaskImage = index;
    state.drawnRect = null;
    state.drawnRectNatural = null;

    const imgData = state.uploadedImages[index];
    const canvas = document.getElementById('mask-canvas');
    const ctx = canvas.getContext('2d');
    const img = new Image();

    img.onload = () => {
        const maxW = 800;
        const maxH = 600;
        let w = img.width;
        let h = img.height;

        if (w > maxW || h > maxH) {
            const ratio = Math.min(maxW / w, maxH / h);
            w = Math.floor(w * ratio);
            h = Math.floor(h * ratio);
        }

        canvas.width = w;
        canvas.height = h;
        canvas.dataset.naturalWidth = img.width;
        canvas.dataset.naturalHeight = img.height;

        ctx.drawImage(img, 0, 0, w, h);
        updateMaskCounter();
        document.getElementById('step4-canvas').classList.remove('hidden');
    };

    img.src = imgData.url;

    let startX, startY;

    function getCanvasCoords(e) {
        const rect = canvas.getBoundingClientRect();
        const scaleX = canvas.width / rect.width;
        const scaleY = canvas.height / rect.height;
        return {
            x: (e.clientX - rect.left) * scaleX,
            y: (e.clientY - rect.top) * scaleY
        };
    }

    canvas.onmousedown = (e) => {
        const pos = getCanvasCoords(e);
        startX = pos.x;
        startY = pos.y;
        state.isDrawing = true;
        state.drawnRect = null;
        state.drawnRectNatural = null;
    };

    canvas.onmousemove = (e) => {
        if (!state.isDrawing) return;
        const pos = getCanvasCoords(e);
        const x = pos.x;
        const y = pos.y;

        ctx.drawImage(img, 0, 0, canvas.width, canvas.height);

        ctx.strokeStyle = '#2563eb';
        ctx.lineWidth = 2;
        ctx.setLineDash([5, 5]);
        ctx.strokeRect(startX, startY, x - startX, y - startY);

        ctx.fillStyle = 'rgba(37, 99, 235, 0.2)';
        ctx.fillRect(startX, startY, x - startX, y - startY);
    };

    canvas.onmouseup = (e) => {
        if (!state.isDrawing) return;
        state.isDrawing = false;

        const pos = getCanvasCoords(e);
        const endX = pos.x;
        const endY = pos.y;

        const x = Math.min(startX, endX);
        const y = Math.min(startY, endY);
        const width = Math.abs(endX - startX);
        const height = Math.abs(endY - startY);

        if (width > 10 && height > 10) {
            state.drawnRect = { x, y, width, height };

            const scaleX = img.width / canvas.width;
            const scaleY = img.height / canvas.height;
            state.drawnRectNatural = {
                x: Math.round(x * scaleX),
                y: Math.round(y * scaleY),
                width: Math.round(width * scaleX),
                height: Math.round(height * scaleY)
            };
        }

        ctx.setLineDash([]);
    };

    canvas.onmouseleave = () => {
        if (state.isDrawing) {
            state.isDrawing = false;
            ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
            if (state.drawnRect) {
                ctx.strokeStyle = '#2563eb';
                ctx.lineWidth = 2;
                ctx.strokeRect(state.drawnRect.x, state.drawnRect.y,
                               state.drawnRect.width, state.drawnRect.height);
            }
        }
    };
}

function toggleInputMode() {
    const mode = document.querySelector('input[name="input_mode"]:checked').value;
    const singleInput = document.getElementById('gen-clean-image');
    const folderInput = document.getElementById('gen-clean-folder');
    const singleDiv = document.getElementById('single-image-input');
    const folderDiv = document.getElementById('folder-input');
    const hint = document.getElementById('input-hint');

    if (mode === 'folder') {
        singleDiv.classList.add('hidden');
        folderDiv.classList.remove('hidden');
        singleInput.value = '';
        hint.textContent = 'Select a folder containing clean (non-defective) product images. All images will be processed.';
    } else {
        singleDiv.classList.remove('hidden');
        folderDiv.classList.add('hidden');
        folderInput.value = '';
        hint.textContent = 'Upload a clean (non-defective) image of the product.';
    }
}

function toggleDDInputMode() {
    const mode = document.querySelector('input[name="dd_input_mode"]:checked').value;
    const singleDiv = document.getElementById('dd-single-image-input');
    const folderDiv = document.getElementById('dd-folder-input');
    const singleInput = document.getElementById('dd-clean-image');
    const folderInput = document.getElementById('dd-clean-folder');
    const hint = document.getElementById('dd-input-hint');

    if (mode === 'folder') {
        singleDiv.classList.add('hidden');
        folderDiv.classList.remove('hidden');
        singleInput.value = '';
        hint.textContent = 'Select a folder containing clean (non-defective) product images. All images will be processed.';
    } else {
        singleDiv.classList.remove('hidden');
        folderDiv.classList.add('hidden');
        folderInput.value = '';
        hint.textContent = 'Upload a clean (non-defective) image of the product.';
    }
}

function clearRect() {
    state.drawnRect = null;
    state.drawnRectNatural = null;
    const canvas = document.getElementById('mask-canvas');
    const ctx = canvas.getContext('2d');
    const img = new Image();
    img.onload = () => {
        ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
    };
    if (state.uploadedImages[state.currentMaskImage]) {
        img.src = state.uploadedImages[state.currentMaskImage].url;
    }
}

function saveMask() {
    if (!state.drawnRectNatural) {
        showAlert('step4-alert', 'Please draw a rectangle on the image first.', 'warning');
        return;
    }

    const imgData = state.uploadedImages[state.currentMaskImage];

    fetch('/api/training/save-mask', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
            image_path: imgData.path,
            rect: state.drawnRectNatural
        })
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            state.savedMasks.add(state.currentMaskImage);
            const badge = document.getElementById(`mask-badge-${state.currentMaskImage}`);
            if (badge) badge.classList.remove('hidden');
            updateMaskCounter();

            const totalSaved = state.savedMasks.size;
            const remaining = state.uploadedImages.length - totalSaved;

            if (totalSaved >= 3) {
                showAlert('step4-alert',
                    `Mask saved! Please click the <strong>Done</strong> button to continue to Step 5 (Remaining unmasked images: ${remaining})`,
                    'success',
                    0);
            } else {
                showAlert('step4-alert', 'Mask saved!', 'success', 1200);
            }

            if (state.currentMaskImage < state.uploadedImages.length - 1) {
                setTimeout(() => nextMaskImage(), 500);
            }
        } else {
            showAlert('step4-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        showAlert('step4-alert', 'Error: ' + err.message, 'danger');
    });
}

function skipMask() {
    if (state.currentMaskImage < state.uploadedImages.length - 1) {
        nextMaskImage();
    }
}

function prevMaskImage() {
    if (state.currentMaskImage > 0) {
        initMaskCanvas(state.currentMaskImage - 1);
    }
}

function nextMaskImage() {
    if (state.currentMaskImage < state.uploadedImages.length - 1) {
        initMaskCanvas(state.currentMaskImage + 1);
    }
}

function updateMaskCounter() {
    const counter = document.getElementById('mask-counter');
    const total = state.uploadedImages.length;
    const masked = state.savedMasks.size;
    const current = state.currentMaskImage !== null ? state.currentMaskImage + 1 : 0;
    const filename = state.uploadedImages[state.currentMaskImage]?.filename || '';
    counter.textContent = `Image ${current} of ${total}: ${filename} | Masks saved: ${masked}/${total}`;
}

function finishMaskStep() {
    if (state.savedMasks.size < 1) {
        showAlert('step4-alert', `You must save masks for at least 1 images. Currently saved: ${state.savedMasks.size}.`, 'warning');
        return;
    }
    document.getElementById('step5').classList.remove('hidden');
    document.getElementById('step5').scrollIntoView({ behavior: 'smooth' });
    document.getElementById('step4-alert').classList.add('hidden');
}

// =============================================================================
// Domain Instructions
// =============================================================================

function setupDomainExamples() {
    const examples = {
        triac: `If the image is a standard TO-220 package (like a Triac):
- The component consists of three metallic pins/leads, a black plastic mold body, a metal mounting tab with a mounting hole, and printed silkscreen text.
- CRITICAL CONSTRAINT: You MUST exclusively propose defects located on the metallic pins/leads. Do NOT propose any defects on the black plastic body, the metal mounting tab, the mounting hole, or the silkscreen text.
- Acceptable pin/lead defects include: pin bending (crooked leads), lead deformation, surface scratches on the metallic pins, or contamination/solder dross on the lead surfaces.
- Use "3 pin" as target_entity rather than specific individual pins.
- The description must clearly detail exactly which pin is affected and how it is deformed or damaged (e.g., "The right pin is bent sharply outwards near the mid-section").
- The defect_coverage_ratio should be 0.4 - 0.60.`,

        xray_pcb: `If the image is an X-ray or CT scan of a PCB/IC package:
- The die appears as a large dark or light rectangular region, usually in the center.
- Solder balls/BGA appear as an array of circular dots.
- Defects in X-ray appear as grayscale anomalies (darker or brighter regions), NOT as circuit traces, lines, or wire patterns.
- A "xray_die" defect should look like: a natural void, crack, or delamination INSIDE the die material — blobby, irregular, or circular grayscale spots.
- A "xray_solder" defect should look like: missing solder, solder bridges, or voids in solder balls.
- NEVER describe the defect as "circuit traces", "wiring", "grid lines", or "PCB patterns". The defect must look like a manufacturing flaw, not an overlaid schematic.
- The defect_coverage_ratio should be 1.0.`,

        vcsel: `If the image is a VCSEL laser diode:
- The component has a metallic can package (TO-can), a glass lens cap on top, and electrode bars.
- Defects should focus on: lens cap cracks, scratches on the glass dome, electrode bar discontinuity, or metallic can surface damage.
- Use "lens cap" for lens-related defects and "electrode bars" for electrode-related defects.
- Avoid targeting the emission aperture directly with large coverage masks.`,

        generic: `If the image shows an electronic component:
- Identify the main functional areas: metallic leads/pins, plastic or ceramic body, and any optical or contact surfaces.
- Propose defects that are realistic manufacturing flaws: scratches, cracks, contamination, bent leads, missing material, or discoloration.
- Use specific terminology for metallic parts (e.g., "metallic pins", "copper leads", "gold-plated contacts") rather than generic terms like "leg" or "metal part".
- The defect must be visually distinct and not confused with normal component features.`
    };

    window.domainExamples = examples;
}

function loadDomainExample() {
    const select = document.getElementById('domain-example-select');
    const key = select.value;
    if (window.domainExamples[key]) {
        document.getElementById('domain-instructions').value = window.domainExamples[key];
    }
}

function saveDomainInstructions() {
    const objClass = document.getElementById('train-object-class').value.trim() || 'default';
    const instructions = document.getElementById('domain-instructions').value.trim();

    fetch('/api/domain-instructions', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
            object_class: objClass,
            instructions: instructions
        })
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            showAlert('step5-alert', data.message, 'success');
            document.getElementById('step6').classList.remove('hidden');
            document.getElementById('step6').scrollIntoView({ behavior: 'smooth' });
        } else {
            showAlert('step5-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        showAlert('step5-alert', 'Error: ' + err.message, 'danger');
    });
}

// =============================================================================
// Training Execution (DefectFill)
// =============================================================================

function startTraining() {
    const objClass = document.getElementById('train-object-class').value.trim();
    const defectType = document.getElementById('train-defect-type').value.trim();

    const config = {
        object_class: objClass,
        defect_type: defectType,
        lora_rank: parseInt(document.getElementById('cfg-lora-rank').value) || 8,
        lora_alpha: parseInt(document.getElementById('cfg-lora-alpha').value) || 16,
        max_train_steps: parseInt(document.getElementById('cfg-max-steps').value) || 1500,
        batch_size: parseInt(document.getElementById('cfg-batch-size').value) || 2,
        gradient_accumulation_steps: parseInt(document.getElementById('cfg-grad-accum').value) || 2,
        lambda_defect: parseFloat(document.getElementById('cfg-lambda-defect').value) || 0.5,
        lambda_obj: parseFloat(document.getElementById('cfg-lambda-obj').value) || 0.2,
        lambda_attn: parseFloat(document.getElementById('cfg-lambda-attn').value) || 0.05,
        alpha: parseFloat(document.getElementById('cfg-alpha').value) || 0.3
    };

    const btn = document.getElementById('btn-start-training');
    btn.disabled = true;
    btn.innerHTML = '<span class="spinner"></span> Starting Training...';

    fetch('/api/training/start', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(config)
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            document.getElementById('training-monitor').classList.remove('hidden');
            document.getElementById('training-monitor').scrollIntoView({
                behavior: 'smooth',
                block: 'start'
            });
            showAlert('step6-alert', 'Training started! Monitor logs below.', 'success');
            startLogPolling();
        } else {
            btn.disabled = false;
            btn.innerHTML = '🚀 Start Training';
            showAlert('step6-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        btn.innerHTML = '🚀 Start Training';
        showAlert('step6-alert', 'Error: ' + err.message, 'danger');
    });
}

function stopTraining() {
    fetch('/api/training/stop', { method: 'POST' })
    .then(r => r.json())
    .then(data => {
        showAlert('step6-alert', data.message || data.error, data.success ? 'warning' : 'danger');
    });
}

function startLogPolling() {
    if (state.trainingPollInterval) clearInterval(state.trainingPollInterval);

    const logViewer = document.getElementById('training-log');
    const statusBadge = document.getElementById('training-status-badge');
    const progressBar = document.getElementById('training-progress');
    const btnStart = document.getElementById('btn-start-training');
    const btnStop = document.getElementById('btn-stop-training');

    btnStop.classList.remove('hidden');

    state.trainingPollInterval = setInterval(() => {
        fetch('/api/training/status')
        .then(r => r.json())
        .then(data => {
            if (!data.success) return;

            const status = data.status;

            statusBadge.className = 'status-badge';
            if (status.running) {
                statusBadge.classList.add('running');
                statusBadge.innerHTML = '<span>●</span> Running';
            } else if (status.error) {
                statusBadge.classList.add('error');
                statusBadge.innerHTML = '<span>●</span> Error';
            } else if (status.completed) {
                statusBadge.classList.add('completed');
                statusBadge.innerHTML = '<span>●</span> Completed';
            } else {
                statusBadge.classList.add('idle');
                statusBadge.innerHTML = '<span>●</span> Idle';
            }

            if (status.log_lines && status.log_lines.length > 0) {
                logViewer.innerHTML = status.log_lines.map(line => {
                    let cls = 'log-line';
                    if (line.includes('Error') || line.includes('FAILED')) cls += ' log-error';
                    else if (line.includes('complete') || line.includes('saved')) cls += ' log-success';
                    else if (line.includes('Warning')) cls += ' log-warn';
                    return `<div class="${cls}">${escapeHtml(line)}</div>`;
                }).join('');
                logViewer.scrollTop = logViewer.scrollHeight;
            }

            const progress = estimateProgress(status.log_lines);
            progressBar.style.width = progress + '%';

            if (!status.running && status.start_time) {
                clearInterval(state.trainingPollInterval);
                btnStart.disabled = false;
                btnStart.innerHTML = '🚀 Start Training';
                btnStop.classList.add('hidden');

                if (status.completed) {
                    showAlert('step6-alert', 'Training completed successfully! You can now generate defects.', 'success');
                    loadCheckpoints();
                } else if (status.error) {
                    showAlert('step6-alert', 'Training failed: ' + status.error, 'danger');
                }
            }
        });
    }, 2000);
}

function estimateProgress(logLines) {
    if (!logLines || logLines.length === 0) return 0;
    const lastLines = logLines.slice(-20).join(' ');
    const match = lastLines.match(/(\d+)\s*\/\s*(\d+)/);
    if (match) {
        const current = parseInt(match[1]);
        const total = parseInt(match[2]);
        if (total > 0) return Math.min(100, (current / total) * 100);
    }
    if (lastLines.includes('complete') || lastLines.includes('saved')) return 100;
    return 0;
}

// =============================================================================
// Generation (DefectFill)
// =============================================================================

function initGenerationForm() {
    document.getElementById('gen-form').addEventListener('submit', (e) => {
        e.preventDefault();
        generateDefect();
    });

    document.getElementById('btn-refresh-checkpoints').addEventListener('click', loadCheckpoints);

    initDetectSimilarProduct();
}

function initDetectSimilarProduct() {
    const detectUploadZone = document.getElementById('detect-upload-zone');
    const detectFileInput = document.getElementById('detect-product-image');
    const btnDetect = document.getElementById('btn-detect-product');
    const preview = document.getElementById('detect-image-preview');

    if (!detectUploadZone || !btnDetect) return;

    detectUploadZone.addEventListener('click', () => detectFileInput.click());

    detectFileInput.addEventListener('change', (e) => {
        if (e.target.files && e.target.files[0]) {
            const file = e.target.files[0];
            preview.innerHTML = `<img src="${URL.createObjectURL(file)}" style="max-width: 100%; border-radius: var(--radius); border: 1px solid var(--gray-200);">`;
            preview.classList.remove('hidden');
        }
    });

    btnDetect.addEventListener('click', () => {
        const files = detectFileInput.files;
        if (!files || files.length === 0) {
            showAlert('detect-product-alert', 'Please upload 1 clean product image first.', 'warning');
            return;
        }

        const formData = new FormData();
        formData.append('image', files[0]);

        btnDetect.disabled = true;
        btnDetect.innerHTML = '<span class="spinner spinner-dark"></span> Detecting similarity...';

        fetch('/api/detect-similar-product', {
            method: 'POST',
            body: formData
        })
        .then(r => r.json())
        .then(data => {
            btnDetect.disabled = false;
            btnDetect.innerHTML = '🔍 Verify Product Model';

            if (data.success) {
                const res = data.result;
                if (res.is_trained) {
                    const select = document.getElementById('gen-object-class');
                    select.value = res.matched_product;
                    select.dispatchEvent(new Event('change'));

                    showAlert(
                        'detect-product-alert',
                        `<strong>✅ Match Found!</strong> Recognized as <strong>${escapeHtml(res.matched_product)}</strong> (Similarity: ${(res.confidence * 100).toFixed(1)}%). Object Class dropdown updated automatically.`,
                        'success'
                    );
                } else {
                    const alertHtml = `
                        <div>
                            <strong>⚠️ Product Not Trained</strong><br>
                            ${escapeHtml(res.message)}<br>
                            <button class="btn btn-sm btn-primary mt-1" onclick="switchToTrainingTab()">
                                🎓 Go to Few-Shot Training
                            </button>
                        </div>
                    `;
                    const alertDiv = document.getElementById('detect-product-alert');
                    alertDiv.className = 'alert alert-warning mt-2';
                    alertDiv.innerHTML = alertHtml;
                    alertDiv.classList.remove('hidden');
                }
            } else {
                showAlert('detect-product-alert', data.error, 'danger');
            }
        })
        .catch(err => {
            btnDetect.disabled = false;
            btnDetect.innerHTML = '🔍 Verify Product Model';
            showAlert('detect-product-alert', 'Detection error: ' + err.message, 'danger');
        });
    });
}

function switchToTrainingTab() {
    const trainingTabBtn = document.querySelector('.tab-btn[data-tab="training"]');
    if (trainingTabBtn) {
        trainingTabBtn.click();
        window.scrollTo({ top: 0, behavior: 'smooth' });
    }
}

function loadCheckpoints() {
    fetch('/api/checkpoints')
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            state.checkpoints = data.checkpoints;
            renderCheckpointSelector();
        }
    })
    .catch(console.error);
}

function renderCheckpointSelector() {
    const select = document.getElementById('gen-object-class');
    const defectSelect = document.getElementById('gen-defect-type');

    const currentObj = select.value;

    select.innerHTML = '<option value="">-- Select Product --</option>';
    Object.keys(state.checkpoints).forEach(objClass => {
        const opt = document.createElement('option');
        opt.value = objClass;
        opt.textContent = objClass;
        select.appendChild(opt);
    });

    if (currentObj && state.checkpoints[currentObj]) {
        select.value = currentObj;
        updateDefectTypeOptions(currentObj);
    }

    select.addEventListener('change', (e) => {
        updateDefectTypeOptions(e.target.value);
    });
}

function updateDefectTypeOptions(objClass) {
    const defectSelect = document.getElementById('gen-defect-type');
    defectSelect.innerHTML = '<option value="">-- Auto (VLM selects) --</option>';

    if (state.checkpoints[objClass]) {
        state.checkpoints[objClass].forEach(dt => {
            const opt = document.createElement('option');
            opt.value = dt;
            opt.textContent = dt;
            defectSelect.appendChild(opt);
        });
    }
}

function generateDefect() {
    const form = document.getElementById('gen-form');
    const formData = new FormData(form);

    const objClass = formData.get('object_class');
    if (!objClass) {
        showAlert('gen-alert', 'Please select a trained product.', 'warning');
        return;
    }

    const btn = document.getElementById('btn-generate');
    btn.disabled = true;
    btn.innerHTML = '<span class="spinner"></span> Generating...';

    fetch('/api/generate', {
        method: 'POST',
        body: formData
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            startGenerationPolling(data.job_id, 'defectfill');
        } else {
            btn.disabled = false;
            btn.innerHTML = '✨ Generate Synthetic Defects';
            showAlert('gen-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        btn.innerHTML = '✨ Generate Synthetic Defects';
        showAlert('gen-alert', 'Generation failed: ' + err.message, 'danger');
    });
}

// =============================================================================
// DefectDiffu
// =============================================================================

function initDefectDiffu() {
    // Load config on init
    loadDDConfig();

    // Save config button
    document.getElementById('btn-save-dd-config').addEventListener('click', saveDDConfig);

    // Generation form
    document.getElementById('dd-gen-form').addEventListener('submit', (e) => {
        e.preventDefault();
        generateDefectDiffu();
    });
}

function loadDDConfig() {
    fetch('/api/defectdiffu/config')
    .then(r => r.json())
    .then(data => {
        if (data.success && data.config) {
            state.ddConfig = data.config;
            document.getElementById('dd-ckpt-path').value = data.config.ckpt_path || '';
            document.getElementById('dd-vae-path').value = data.config.vae_path || '';
            document.getElementById('dd-vlm-model').value = data.config.vlm_model || 'gemma3:12b';
        }
    })
    .catch(console.error);
}

function saveDDConfig() {
    const ckptPath = document.getElementById('dd-ckpt-path').value.trim();
    const vaePath = document.getElementById('dd-vae-path').value.trim();
    const vlmModel = document.getElementById('dd-vlm-model').value.trim() || 'gemma3:12b';

    if (!ckptPath || !vaePath) {
        showAlert('dd-config-alert', 'Both checkpoint path and VAE path are required.', 'warning');
        return;
    }

    fetch('/api/defectdiffu/config', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
            ckpt_path: ckptPath,
            vae_path: vaePath,
            vlm_model: vlmModel
        })
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            state.ddConfig = { ckpt_path: ckptPath, vae_path: vaePath, vlm_model: vlmModel };
            showAlert('dd-config-alert', data.message, 'success');
        } else {
            showAlert('dd-config-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        showAlert('dd-config-alert', 'Error saving config: ' + err.message, 'danger');
    });
}

function generateDefectDiffu() {
    const form = document.getElementById('dd-gen-form');
    const formData = new FormData(form);

    const productDesc = formData.get('product_desc');
    if (!productDesc) {
        showAlert('dd-gen-alert', 'Product description is required.', 'warning');
        return;
    }

    // Check config is set
    if (!state.ddConfig.ckpt_path || !state.ddConfig.vae_path) {
        showAlert('dd-gen-alert', 'Please configure DefectDiffu model paths first.', 'warning');
        document.getElementById('dd-config-card').scrollIntoView({ behavior: 'smooth' });
        return;
    }

    const btn = document.getElementById('btn-dd-generate');
    btn.disabled = true;
    btn.innerHTML = '<span class="spinner"></span> Generating...';

    fetch('/api/defectdiffu/generate', {
        method: 'POST',
        body: formData
    })
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            startGenerationPolling(data.job_id, 'defectdiffu');
        } else {
            btn.disabled = false;
            btn.innerHTML = '✨ Generate with DefectDiffu';
            showAlert('dd-gen-alert', data.error, 'danger');
        }
    })
    .catch(err => {
        btn.disabled = false;
        btn.innerHTML = '✨ Generate with DefectDiffu';
        showAlert('dd-gen-alert', 'Generation failed: ' + err.message, 'danger');
    });
}

// =============================================================================
// Generation Polling (Unified for both models)
// =============================================================================

function startGenerationPolling(jobId, mode) {
    const btn = mode === 'defectfill'
        ? document.getElementById('btn-generate')
        : document.getElementById('btn-dd-generate');
    const resultsContainer = mode === 'defectfill'
        ? document.getElementById('gen-results')
        : document.getElementById('dd-gen-results');
    const summaryDiv = mode === 'defectfill'
        ? document.getElementById('gen-summary')
        : document.getElementById('dd-gen-summary');
    const alertId = mode === 'defectfill' ? 'gen-alert' : 'dd-gen-alert';

    resultsContainer.innerHTML = `
        <div class="card" style="padding: 1.5rem;">
            <h3 style="margin-bottom: 1rem;">⏳ Generating Synthetic Defects...</h3>
            <div class="progress-bar mb-2">
                <div class="fill" id="gen-progress-bar-${mode}" style="width: 0%"></div>
            </div>
            <div style="display: flex; justify-content: space-between; align-items: center;">
                <span id="gen-step-text-${mode}" style="font-weight: 500; color: var(--primary);">Initializing...</span>
                <span id="gen-defect-counter-${mode}" style="font-size: 0.875rem; color: var(--gray-500);"></span>
            </div>
            <div id="gen-log-viewer-${mode}" class="log-viewer mt-2" style="max-height: 200px; font-size: 0.8rem;">
                <div class="log-line text-gray">Waiting for generation to start...</div>
            </div>
        </div>
    `;

    const pollInterval = setInterval(() => {
        fetch(`/api/generation-status/${jobId}`)
        .then(r => r.json())
        .then(data => {
            if (!data.success) {
                clearInterval(pollInterval);
                return;
            }

            const job = data.job;
            let progress = 0;
            if (job.unit_total && job.unit_current !== undefined) {
                progress = Math.min(100, (job.unit_current / job.unit_total) * 100);
            } else {
                progress = Math.min(100, (job.step / job.total_steps) * 100);
            }

            document.getElementById(`gen-progress-bar-${mode}`).style.width = progress + '%';
            document.getElementById(`gen-step-text-${mode}`).textContent =
                `Step ${job.step}/${job.total_steps}: ${job.step_text}`;

            const defectCounter = document.getElementById(`gen-defect-counter-${mode}`);
            let counterText = '';
            if (job.image_total > 0) {
                counterText += `Image ${job.image_current}/${job.image_total}`;
            }
            if (job.defect_total > 0) {
                counterText += (counterText ? ' | ' : '') + `Defect ${job.defect_current}/${job.defect_total}`;
            }
            defectCounter.textContent = counterText;

            const logViewer = document.getElementById(`gen-log-viewer-${mode}`);
            const line = document.createElement('div');
            line.className = 'log-line';
            line.textContent = `[${new Date().toLocaleTimeString()}] ${job.step_text}`;
            logViewer.appendChild(line);
            logViewer.scrollTop = logViewer.scrollHeight;

            if (job.status === 'completed') {
                clearInterval(pollInterval);
                btn.disabled = false;
                btn.innerHTML = mode === 'defectfill'
                    ? '✨ Generate Synthetic Defects'
                    : '✨ Generate with DefectDiffu';
                renderGenerationResults(job.result, mode);
                showAlert(alertId, `Generated ${job.result.output_images?.length || 0} images in ${formatElapsedTime(job.result.elapsed_time)}`, 'success');
            } else if (job.status === 'error') {
                clearInterval(pollInterval);
                btn.disabled = false;
                btn.innerHTML = mode === 'defectfill'
                    ? '✨ Generate Synthetic Defects'
                    : '✨ Generate with DefectDiffu';
                showAlert(alertId, 'Generation failed: ' + job.error, 'danger');
            }
        })
        .catch(console.error);
    }, 10000);
}

function renderGenerationResults(data, mode = 'defectfill') {
    const container = mode === 'defectfill'
        ? document.getElementById('gen-results')
        : document.getElementById('dd-gen-results');
    const summary = mode === 'defectfill'
        ? document.getElementById('gen-summary')
        : document.getElementById('dd-gen-summary');

    container.innerHTML = '';

    const grouped = data.grouped_results;
    let totalImages = 0;
    if (grouped) {
        grouped.forEach(g => {
            totalImages += g.defects.length * 4;
        });
    }

    if (grouped && grouped.length > 0) {
        summary.innerHTML = `
            <div class="alert alert-success">
                <strong>Generation Complete</strong><br>
                Experiment: ${data.experiment_id || 'N/A'} |
                Product: ${data.product_type || 'N/A'} |
                Time: ${formatElapsedTime(data.elapsed_time)} |
                ${totalImages} image(s) generated
            </div>
        `;

        grouped.forEach((group, groupIdx) => {
            const groupId = `result-group-${mode}-${groupIdx}`;

            const groupEl = document.createElement('div');
            groupEl.className = 'result-group';

            const header = document.createElement('div');
            header.className = 'result-group-header';
            header.onclick = () => toggleResultGroup(groupId);
            header.innerHTML = `
                <span class="group-title">Source ${groupIdx + 1}: ${escapeHtml(group.source_name)}</span>
                <span class="group-meta">${group.defect_count} defect(s)</span>
                <span class="toggle-icon" id="icon-${groupId}">▶</span>
            `;

            const content = document.createElement('div');
            content.className = 'result-group-content';
            content.id = groupId;

            group.defects.forEach((defect, defectIdx) => {
                const defectId = `${groupId}-defect-${defectIdx}`;

                const defectGroup = document.createElement('div');
                defectGroup.className = 'result-defect-group';

                const defectHeader = document.createElement('div');
                defectHeader.className = 'result-defect-header';
                defectHeader.onclick = (e) => {
                    e.stopPropagation();
                    toggleResultGroup(defectId);
                };
                defectHeader.innerHTML = `
                    <span class="defect-title">${escapeHtml(defect.defect_label)}</span>
                    <span class="defect-meta">${escapeHtml(defect.defect_type)}</span>
                    <span class="toggle-icon" id="icon-${defectId}">▶</span>
                `;

                const defectContent = document.createElement('div');
                defectContent.className = 'result-defect-content';
                defectContent.id = defectId;

                const trioGrid = document.createElement('div');
                trioGrid.className = 'result-trio-grid';

                const imageTypes = [
                    { key: 'clean_image', label: 'Clean' },
                    { key: 'mask_image', label: 'Mask' },
                    { key: 'raw_patch_image', label: 'Raw Patch' },
                    { key: 'output_image', label: 'Output' }
                ];

                imageTypes.forEach(type => {
                    const imgData = defect[type.key];
                    if (imgData && imgData.url) {
                        const card = document.createElement('div');
                        card.className = 'result-trio-card';
                        card.innerHTML = `
                            <div class="trio-label">${type.label}</div>
                            <img src="${imgData.url}" alt="${imgData.filename}" onclick="previewImage('${imgData.url}')">
                            <div class="trio-filename">${escapeHtml(imgData.filename)}</div>
                            <a href="${imgData.url}" download class="btn btn-sm btn-secondary mt-1">Download</a>
                        `;
                        trioGrid.appendChild(card);
                    }
                });

                defectContent.appendChild(trioGrid);
                defectGroup.appendChild(defectHeader);
                defectGroup.appendChild(defectContent);
                content.appendChild(defectGroup);
            });

            groupEl.appendChild(header);
            groupEl.appendChild(content);
            container.appendChild(groupEl);
        });
    } else if (data.output_images && data.output_images.length > 0) {
        summary.innerHTML = `
            <div class="alert alert-success">
                <strong>Generation Complete</strong><br>
                Experiment: ${data.experiment_id || 'N/A'} |
                Product: ${data.product_type || 'N/A'} |
                Time: ${formatElapsedTime(data.elapsed_time)}
            </div>
        `;

        const grid = document.createElement('div');
        grid.className = 'result-grid';

        data.output_images.forEach(img => {
            const card = document.createElement('div');
            card.className = 'result-card';
            card.innerHTML = `
                <img src="${img.url}" alt="${img.filename}" onclick="previewImage('${img.url}')">
                <div class="info">
                    <h4>${img.defect_type}</h4>
                    <p>${img.filename}</p>
                    <a href="${img.url}" download class="btn btn-sm btn-secondary mt-1">Download</a>
                </div>
            `;
            grid.appendChild(card);
        });

        container.appendChild(grid);

    } else {
        summary.innerHTML = `
            <div class="alert alert-warning">
                No images were generated. Check the defect type and product description.
            </div>
        `;
    }
}

function toggleResultGroup(groupId) {
    const content = document.getElementById(groupId);
    const icon = document.getElementById(`icon-${groupId}`);
    const header = content.previousElementSibling;

    if (content.classList.contains('expanded')) {
        content.classList.remove('expanded');
        icon.textContent = '▶';
        header.classList.remove('active');
    } else {
        content.classList.add('expanded');
        icon.textContent = '▼';
        header.classList.add('active');
    }
}


// =============================================================================
// Utilities
// =============================================================================

function showAlert(elementId, message, type, autoHideMs = null) {
    const el = document.getElementById(elementId);
    if (!el) return;

    if (el._hideTimeout) {
        clearTimeout(el._hideTimeout);
        el._hideTimeout = null;
    }

    el.className = `alert alert-${type}`;
    el.innerHTML = message;
    el.classList.remove('hidden');

    if (type === 'success') {
        const delay = autoHideMs !== null ? autoHideMs : 0;
        if (delay > 0) {
            el._hideTimeout = setTimeout(() => {
                el.classList.add('hidden');
            }, delay);
        }
    }
}

function escapeHtml(text) {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
}

function previewImage(url) {
    const overlay = document.createElement('div');
    overlay.className = 'modal-overlay';
    overlay.innerHTML = `
        <div class="modal-content">
            <button class="modal-close" onclick="this.parentElement.parentElement.remove()">×</button>
            <img src="${url}" alt="Preview">
        </div>
    `;
    overlay.addEventListener('click', (e) => {
        if (e.target === overlay) overlay.remove();
    });
    document.body.appendChild(overlay);
}

function loadSystemInfo() {
    fetch('/api/system/info')
    .then(r => r.json())
    .then(data => {
        if (data.success) {
            const info = data.info;
            const badge = document.getElementById('cuda-badge');
            if (info.cuda_available) {
                badge.textContent = `CUDA: ${info.cuda_device_name}`;
                badge.style.background = 'var(--success)';
            } else {
                badge.textContent = 'CPU Mode';
                badge.style.background = 'var(--warning)';
            }
        }
    })
    .catch(console.error);
}

function formatElapsedTime(seconds) {
    if (!seconds && seconds !== 0) return '0s';
    const mins = Math.floor(seconds / 60);
    const secs = (seconds % 60).toFixed(1);
    if (mins > 0) {
        return `${mins}m ${secs}s`;
    }
    return `${secs}s`;
}