document.addEventListener('DOMContentLoaded', function() {
    // Add listeners for radio buttons
    const radioButtons = document.querySelectorAll('input[name="inputMethod"]');
    radioButtons.forEach(radio => {
        radio.addEventListener('change', toggleInputMethod);
    });

    // Initialize the search and labels functionality
    initializeSearchAndLabels();
});

function toggleInputMethod() {
    const localInputGroup = document.getElementById('localInputGroup');
    const urlInputGroup = document.getElementById('urlInputGroup');
    const selectedMethod = document.querySelector('input[name="inputMethod"]:checked').value;

    if (selectedMethod === 'local') {
        localInputGroup.style.display = 'block';
        urlInputGroup.style.display = 'none';
        // Clear URL input
        document.getElementById('urlInput').value = '';
    } else {
        localInputGroup.style.display = 'none';
        urlInputGroup.style.display = 'block';
        // Clear file input
        document.getElementById('fileInput').value = '';
    }
    
    // Clear preview and results
    clearPreviewAndResults();
}

function clearPreviewAndResults() {
    const imagePreview = document.getElementById('imagePreview');
    const resultsDiv = document.getElementById('results');
    
    imagePreview.src = '';
    resultsDiv.innerHTML = '';
}

// Auto-preview when file input changes
document.getElementById('fileInput').addEventListener('change', previewImage);
document.getElementById('urlInput').addEventListener('input', debounce(previewImage, 500));

function debounce(func, wait) {
    let timeout;
    return function executedFunction(...args) {
        const later = () => {
            clearTimeout(timeout);
            func(...args);
        };
        clearTimeout(timeout);
        timeout = setTimeout(later, wait);
    };
}

function previewImage() {
    const fileInput = document.getElementById('fileInput');
    const urlInput = document.getElementById('urlInput');
    const imagePreview = document.getElementById('imagePreview');
    const defaultPreviewMessage = document.getElementById('defaultPreviewMessage');

    if (fileInput.files && fileInput.files[0]) {
        const reader = new FileReader();
        reader.onload = function(e) {
            imagePreview.src = e.target.result;
            imagePreview.style.display = 'block'; // Show the image
            defaultPreviewMessage.style.display = 'none'; // Hide the default message
        };
        reader.readAsDataURL(fileInput.files[0]);
        urlInput.value = ''; // Clear URL input
    } else if (urlInput.value) {
        imagePreview.src = urlInput.value;
        imagePreview.style.display = 'block'; // Show the image
        defaultPreviewMessage.style.display = 'none'; // Hide the default message
        fileInput.value = ''; // Clear file input
    } else {
        // If no image is loaded, show the default message
        imagePreview.style.display = 'none'; // Hide the image
        defaultPreviewMessage.style.display = 'block'; // Show the default message
    }
}

async function getBase64FromUrl(url) {
    const response = await fetch(url);
    const blob = await response.blob();
    return new Promise((resolve, reject) => {
        const reader = new FileReader();
        reader.onload = () => {
            const base64String = reader.result.split(',')[1];
            resolve(base64String);
        };
        reader.onerror = reject;
        reader.readAsDataURL(blob);
    });
}

async function analyzeImage() {
    // Clear previous results first
    document.getElementById('results').innerHTML = '';
    
    const spinner = document.getElementById('spinner');
    const analyzeBtn = document.getElementById('analyzeBtn');
    const resultsDiv = document.getElementById('results');
    
    spinner.style.display = 'block';
    analyzeBtn.disabled = true;
    analyzeBtn.textContent = 'Analyzing...';

    const fileInput = document.getElementById('fileInput');
    const urlInput = document.getElementById('urlInput');
    let base64Image;

    try {
        // Get selected labels from both default checkboxes and selected labels container
        const defaultLabels = Array.from(document.querySelectorAll('.default-labels input[type="checkbox"]:checked'))
            .map(checkbox => checkbox.value);
            
        const selectedLabels = Array.from(document.querySelectorAll('#selectedLabels .selected-label'))
            .map(label => label.dataset.label);

        // Combine both sets of labels
        const allSelectedLabels = [...defaultLabels, ...selectedLabels];

        if (allSelectedLabels.length === 0) {
            throw new Error('Please select at least one category');
        }

        if (fileInput.files && fileInput.files[0]) {
            // Convert file to base64
            const reader = new FileReader();
            base64Image = await new Promise((resolve) => {
                reader.onload = (e) => {
                    resolve(e.target.result.split(',')[1]);
                };
                reader.readAsDataURL(fileInput.files[0]);
            });
        } else if (urlInput.value) {
            // Convert URL to base64
            base64Image = await getBase64FromUrl(urlInput.value);
        } else {
            throw new Error('Please select a file or enter a URL.');
        }

        const payload = {
            images: [base64Image],
            labels: allSelectedLabels,
            multilabel: false
        };

        const response = await fetch('http://localhost:8000/predict', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json'
            },
            body: JSON.stringify(payload)
        });

        const data = await response.json();
        displayResults(data.predictions);
    } catch (error) {
        console.error('Error:', error);
        resultsDiv.innerHTML = `<p class="error-message">Error: ${error.message}</p>`;
    } finally {
        spinner.style.display = 'none';
        analyzeBtn.disabled = false;
        analyzeBtn.textContent = 'Analyze';
    }
}

function displayResults(predictions) {
    const resultsDiv = document.getElementById('results');
    resultsDiv.innerHTML = '';

    if (!predictions || predictions.length === 0) {
        resultsDiv.innerHTML = '<p>No predictions available.</p>';
        return;
    }

    predictions.forEach(prediction => {
        resultsDiv.innerHTML += '<h2>Results:</h2>';
        for (const [label, probability] of Object.entries(prediction)) {
            const percentage = (probability * 100).toFixed(1);
            resultsDiv.innerHTML += `
                <div class="prediction-result">
                    <strong>${label}:</strong>
                    <div class="progress-bar">
                        <div class="progress" style="width: ${percentage}%"></div>
                    </div>
                    <span>${percentage}%</span>
                </div>`;
        }
    });
}

// Define your labels with categories
const medicalLabels = {
    'Others': ['Effusion', 'Edema', 'Scar Tissue', 'Calcification'],
    'Pulmonology': ['Pneumonia', 'Asthma', 'COPD', 'Tuberculosis', 'Pulmonary Fibrosis', 'Pleural Effusion', 'Pulmonary Edema', 'Lung Nodule', 'Atelectasis', 'Pneumothorax'],
    'Oncology': ['Tumor', 'Breast Cancer', 'Lung Cancer', 'Prostate Cancer', 'Leukemia', 'Lymphoma', 'Melanoma', 'Colorectal Cancer', 'Glioma', 'Metastasis'],
    'Orthopedics': ['Fracture', 'Arthritis', 'Osteoporosis', 'Scoliosis', 'Tendonitis', 'Joint Effusion', 'Disc Herniation', 'Osteomyelitis', 'Bursitis', 'Bone Lesion'],
    'Cardiology': ['Myocardial Infarction', 'Arrhythmia', 'Heart Failure', 'Cardiomegaly', 'Aortic Aneurysm', 'Valvular Heart Disease', 'Coronary Artery Disease', 'Pericardial Effusion', 'Pulmonary Embolism'],
    'Dermatology': ['Urtikaria', 'Akne', 'Eczema', 'Psoriasis', 'Melanoma', 'Basal Cell Carcinoma', 'Squamous Cell Carcinoma', 'Skin Ulcer', 'Rash'],
    'Gastroenterology': ['Cirrhosis', 'Hepatitis', 'Ulcer', 'Gastric Cancer', 'Polyp', 'Pancreatitis', 'Cholecystitis', 'Colitis', 'Crohn’s Disease'],
    'Neurology': ['Stroke', 'Multiple Sclerosis', 'Parkinson’s Disease', 'Alzheimer’s Disease', 'Epilepsy', 'Brain Tumor', 'Hydrocephalus', 'Meningitis', 'Intracranial Hemorrhage'],
    'Endocrinology': ['Diabetes', 'Thyroid Nodule', 'Goiter', 'Adrenal Tumor', 'Pituitary Adenoma', 'Hyperthyroidism', 'Hypothyroidism', 'Parathyroid Hyperplasia'],
    'Hematology': ['Anemia', 'Thrombocytopenia', 'Leukemia', 'Lymphoma', 'Hemophilia', 'Polycythemia', 'Sickle Cell Disease'],
    'Urology': ['Kidney Stone', 'Bladder Cancer', 'Prostate Cancer', 'Urinary Tract Infection', 'Renal Cyst', 'Hydronephrosis'],
    'Ophthalmology': ['Glaucoma', 'Cataract', 'Retinal Detachment', 'Macular Degeneration', 'Diabetic Retinopathy', 'Conjunctivitis'],
    'Gynecology': ['Ovarian Cyst', 'Fibroids', 'Endometriosis', 'Breast Cancer', 'Polycystic Ovary Syndrome', 'Cervical Cancer'],
    'Rheumatology': ['Rheumatoid Arthritis', 'Lupus', 'Scleroderma', 'Ankylosing Spondylitis', 'Gout', 'Sjogren’s Syndrome']
    // Add more categories and labels as needed
};

function initializeSearchAndLabels() {
    // First, remove any existing search containers to prevent duplicates
    const existingSearchContainers = document.querySelectorAll('.search-container');
    existingSearchContainers.forEach(container => container.remove());

    // Create search container with show all button
    const searchContainer = document.createElement('div');
    searchContainer.className = 'search-container';
    searchContainer.innerHTML = `
        <div class="search-wrapper">
            <input type="text" id="labelSearch" class="search-input" placeholder="Search labels...">
            <button id="showAllLabels" class="show-all-btn">Show All</button>
        </div>
        <div id="labelDropdown" class="label-dropdown" style="display: none;"></div>
    `;

    // Insert the search container after the default labels
    const defaultLabels = document.querySelector('.default-labels');
    if (defaultLabels) {
        defaultLabels.parentNode.insertBefore(searchContainer, defaultLabels.nextElementSibling);
    }

    // Set up event listeners
    const searchInput = document.getElementById('labelSearch');
    const labelDropdown = document.getElementById('labelDropdown');
    const showAllButton = document.getElementById('showAllLabels');

    // Show dropdown on focus
    searchInput.addEventListener('focus', () => {
        const searchTerm = searchInput.value.trim().toLowerCase();
        showFilteredResults(searchTerm);
    });

    // Handle input changes
    searchInput.addEventListener('input', (e) => {
        const searchTerm = e.target.value.trim().toLowerCase();
        if (searchTerm) {
            showFilteredResults(searchTerm);
        } else {
            labelDropdown.style.display = 'none';
        }
    });

    // Show all button click handler
    showAllButton.addEventListener('click', function(e) {
        e.stopPropagation(); // Prevent event from bubbling up
        labelDropdown.style.display = 'block';
        showFilteredResults('');
        searchInput.value = '';
    });

    // Close dropdown when clicking outside
    document.addEventListener('click', (e) => {
        if (!searchInput.contains(e.target) && 
            !labelDropdown.contains(e.target) && 
            !showAllButton.contains(e.target)) {
            labelDropdown.style.display = 'none';
        }
    });
}

function showFilteredResults(searchTerm) {
    const labelDropdown = document.getElementById('labelDropdown');
    labelDropdown.innerHTML = '';
    let hasResults = false;

    // Sort categories alphabetically for better organization
    const sortedCategories = Object.entries(medicalLabels).sort((a, b) => a[0].localeCompare(b[0]));

    // Loop through all medical labels
    sortedCategories.forEach(([category, labels]) => {
        // Always show all labels when searchTerm is empty (Show All button clicked)
        // Otherwise, filter based on search term
        const matchingLabels = searchTerm === '' ? 
            labels.sort() : // Sort labels alphabetically
            labels.filter(label => label.toLowerCase().includes(searchTerm.toLowerCase()));

        if (matchingLabels.length > 0) {
            hasResults = true;
            
            // Add category header
            const categoryHeader = document.createElement('div');
            categoryHeader.className = 'dropdown-category-header';
            categoryHeader.textContent = category;
            labelDropdown.appendChild(categoryHeader);

            // Add labels
            matchingLabels.forEach(label => {
                const option = document.createElement('div');
                option.className = 'dropdown-option';
                
                const checkbox = document.createElement('input');
                checkbox.type = 'checkbox';
                checkbox.id = `${category}-${label}`;
                checkbox.checked = isLabelSelected(category, label);
                checkbox.addEventListener('change', () => updateSelectedLabels(category, label));

                const labelText = document.createElement('span');
                labelText.className = 'label-text';
                labelText.textContent = label;

                const categoryBadge = document.createElement('span');
                categoryBadge.className = 'category-badge';
                categoryBadge.textContent = category;

                option.appendChild(checkbox);
                option.appendChild(labelText);
                option.appendChild(categoryBadge);
                labelDropdown.appendChild(option);
            });
        }
    });

    labelDropdown.style.display = hasResults || searchTerm === '' ? 'block' : 'none';
}

function isLabelSelected(category, label) {
    const selectedLabelsContainer = document.getElementById('selectedLabels');
    return !!selectedLabelsContainer.querySelector(
        `[data-category="${category}"][data-label="${label}"]`
    );
}

function updateSelectedLabels(category, label) {
    const selectedLabelsContainer = document.getElementById('selectedLabels');
    const checkbox = document.getElementById(`${category}-${label}`);
    
    if (checkbox.checked) {
        const labelElement = document.createElement('span');
        labelElement.className = 'selected-label';
        labelElement.textContent = `${label}`;
        labelElement.dataset.category = category;
        labelElement.dataset.label = label;
        
        const categoryBadge = document.createElement('span');
        categoryBadge.className = 'category-badge';
        categoryBadge.textContent = category;
        
        const removeButton = document.createElement('span');
        removeButton.className = 'remove-label';
        removeButton.innerHTML = '×';
        removeButton.onclick = () => {
            labelElement.remove();
            checkbox.checked = false;
        };
        
        labelElement.appendChild(categoryBadge);
        labelElement.appendChild(removeButton);
        selectedLabelsContainer.appendChild(labelElement);
    } else {
        const existingLabel = selectedLabelsContainer.querySelector(
            `[data-category="${category}"][data-label="${label}"]`
        );
        if (existingLabel) {
            existingLabel.remove();
        }
    }
}

// Initialize when DOM is loaded
document.addEventListener('DOMContentLoaded', initializeSearchAndLabels);