let uploadedFile = null;

// DOM elements
const uploadArea = document.getElementById('uploadArea');
const fileInput = document.getElementById('fileInput');
const imageSection = document.getElementById('imageSection');
const uploadedImage = document.getElementById('uploadedImage');
const predictBtn = document.getElementById('predictBtn');
const resultsSection = document.getElementById('resultsSection');
const hotDogProb = document.getElementById('hotDogProb');
const notHotDogProb = document.getElementById('notHotDogProb');
const loading = document.getElementById('loading');

// Drag and drop event listeners
uploadArea.addEventListener('dragover', handleDragOver);
uploadArea.addEventListener('dragleave', handleDragLeave);
uploadArea.addEventListener('drop', handleDrop);
uploadArea.addEventListener('click', () => fileInput.click());

// File input change listener
fileInput.addEventListener('change', handleFileSelect);

function handleDragOver(e) {
    e.preventDefault();
    uploadArea.classList.add('dragover');
}

function handleDragLeave(e) {
    e.preventDefault();
    uploadArea.classList.remove('dragover');
}

function handleDrop(e) {
    e.preventDefault();
    uploadArea.classList.remove('dragover');
    
    const files = e.dataTransfer.files;
    if (files.length > 0) {
        handleFile(files[0]);
    }
}

function handleFileSelect(e) {
    const files = e.target.files;
    if (files.length > 0) {
        handleFile(files[0]);
    }
}

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

    uploadedFile = file;
    
    // Display the image
    const reader = new FileReader();
    reader.onload = function(e) {
        uploadedImage.src = e.target.result;
        imageSection.style.display = 'block';
        predictBtn.disabled = false;
        
        // Hide results if they were previously shown
        resultsSection.style.display = 'none';
    };
    reader.readAsDataURL(file);
}

async function predictImage() {
    if (!uploadedFile) {
        alert('Please upload an image first.');
        return;
    }

    // Show loading state
    loading.style.display = 'block';
    predictBtn.disabled = true;
    resultsSection.style.display = 'none';

    try {
        // Create FormData to send the image
        const formData = new FormData();
        formData.append('file', uploadedFile);

        // Make API call to the predict endpoint
        const response = await fetch('http://localhost:7860/predict', {
            method: 'POST',
            body: formData
        });

        if (!response.ok) {
            throw new Error(`HTTP error! status: ${response.status}`);
        }

        const result = await response.json();
        
        // Display results
        displayResults(result);
        
    } catch (error) {
        console.error('Error making prediction:', error);
        alert('Error making prediction. Please try again.');
    } finally {
        // Hide loading state
        loading.style.display = 'none';
        predictBtn.disabled = false;
    }
}

function displayResults(result) {
    // Extract probabilities from the response
    // Assuming the API returns something like: { "hot_dog": 0.85, "not_hot_dog": 0.15 }
    // or { "predictions": [{"label": "hot_dog", "probability": 0.85}, {"label": "not_hot_dog", "probability": 0.15}] }
    
    let hotDogProbability = 0;
    let notHotDogProbability = 0;

    // Handle different possible response formats
    if (result.hot_dog !== undefined && result.not_hot_dog !== undefined) {
        hotDogProbability = result.hot_dog;
        notHotDogProbability = result.not_hot_dog;
    } else if (result.predictions) {
        result.predictions.forEach(pred => {
            if (pred.label === 'hot_dog') {
                hotDogProbability = pred.probability;
            } else if (pred.label === 'not_hot_dog') {
                notHotDogProbability = pred.probability;
            }
        });
    } else if (result.probabilities) {
        // Handle array format
        hotDogProbability = result.probabilities[0] || 0;
        notHotDogProbability = result.probabilities[1] || 0;
    } else {
        // Fallback: try to extract from any available format
        const keys = Object.keys(result);
        if (keys.length >= 2) {
            const values = Object.values(result);
            hotDogProbability = values[0] || 0;
            notHotDogProbability = values[1] || 0;
        }
    }

    // Format probabilities as percentages
    hotDogProb.textContent = `${(hotDogProbability * 100).toFixed(2)}%`;
    notHotDogProb.textContent = `${(notHotDogProbability * 100).toFixed(2)}%`;

    // Show results section
    resultsSection.style.display = 'block';
}

// Initialize button state
predictBtn.disabled = true;