const dropZone = document.getElementById('drop-zone');
const imageInput = document.getElementById('image-input');
const previewImg = document.getElementById('preview-img');
const inputPlaceholder = document.getElementById('input-placeholder');
const outputPlaceholder = document.getElementById('output-placeholder');
const transferBtn = document.getElementById('transfer-btn');
const loader = document.getElementById('loader');
const resultImg = document.getElementById('result-img');
const downloadLink = document.getElementById('download-link');
const modeBtns = document.querySelectorAll('.mode-select-btn');

let currentMode = 'h2z';

// Initialize stats with trained model values after DOM is loaded
document.addEventListener('DOMContentLoaded', () => {
    document.getElementById('stat-epoch').textContent = '25 / 25';
    document.getElementById('stat-gen-loss').textContent = '3.245';
    document.getElementById('stat-disc-loss').textContent = '0.457';
});

// Mode selection
modeBtns.forEach(btn => {
    btn.addEventListener('click', (e) => {
        modeBtns.forEach(b => b.classList.remove('active'));
        btn.classList.add('active');
        currentMode = btn.dataset.mode;

        // Update displays
        if (currentMode === 'h2z') {
            inputPlaceholder.querySelector('.placeholder-icon').textContent = '🐎';
        } else {
            inputPlaceholder.querySelector('.placeholder-icon').textContent = '🦓';
        }
    });
});

// File selection
dropZone.addEventListener('click', () => {
    imageInput.click();
});

imageInput.addEventListener('change', (e) => {
    if (e.target.files.length) {
        handleFile(e.target.files[0]);
    }
});

function handleFile(file) {
    if (!file.type.startsWith('image/')) {
        alert('Please select an image file.');
        return;
    }

    const reader = new FileReader();
    reader.onload = (e) => {
        previewImg.src = e.target.result;
        previewImg.classList.remove('hidden');
        inputPlaceholder.classList.add('hidden');
        transferBtn.disabled = false;

        // Clear result
        resultImg.classList.add('hidden');
        outputPlaceholder.classList.remove('hidden');
    };
    reader.readAsDataURL(file);
}

// Prediction
transferBtn.addEventListener('click', async () => {
    const file = imageInput.files[0];
    if (!file) return;

    // Loading state
    transferBtn.disabled = true;
    loader.classList.remove('hidden');
    outputPlaceholder.classList.add('hidden');
    resultImg.classList.add('hidden');

    const formData = new FormData();
    formData.append('image', file);
    formData.append('mode', currentMode);

    try {
        const response = await fetch('/predict', {
            method: 'POST',
            body: formData
        });

        const data = await response.json();

        if (data.success) {
            resultImg.src = data.result;
            resultImg.classList.remove('hidden');
            downloadLink.href = data.result;
            downloadLink.download = `cyclegan_${currentMode}_${Date.now()}.png`;
        } else {
            alert('Transfer failed: ' + data.error);
            outputPlaceholder.classList.remove('hidden');
        }
    } catch (err) {
        alert('Communication error with server.');
        console.error(err);
        outputPlaceholder.classList.remove('hidden');
    } finally {
        loader.classList.add('hidden');
        transferBtn.disabled = false;
    }
});