import { AutoProcessor, CLIPVisionModelWithProjection, RawImage, env } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.6.0';

// Since we will download the model from the Hugging Face Hub, we can skip the local model check
env.allowLocalModels = false;

// Reference the elements that we will need
const statusText = document.getElementById('status-text');
const fileUpload = document.getElementById('file-upload');
const dropZone = document.getElementById('drop-zone');
const imagePreview1 = document.getElementById('image-preview-1');
const imagePreview2 = document.getElementById('image-preview-2');
const meterContainer = document.getElementById('meter-container');
const spinner = document.querySelector('.spinner');
const showGraphBtn = document.getElementById('show-graph-btn');
const graphModal = document.getElementById('graph-modal');
const closeModalBtn = document.querySelector('.close-button');
const resetZoomBtn = document.getElementById('reset-zoom-btn');
const graphContainerModal = document.getElementById('graph-container-modal');

// Load processor and vision model for more direct embedding control
statusText.textContent = 'Loading model...';
spinner.style.display = 'block';
const processor = await AutoProcessor.from_pretrained('Xenova/clip-vit-base-patch16');
const vision_model = await CLIPVisionModelWithProjection.from_pretrained('Xenova/clip-vit-base-patch16');
statusText.textContent = 'Ready';
spinner.style.display = 'none';

let imageSrc1 = null;
let imageSrc2 = null;
let lastEmbeds = null;

// Initial setup of upload placeholders
clearUploads();


// Prevent default drag behaviors
['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
    dropZone.addEventListener(eventName, preventDefaults, false);
    document.body.addEventListener(eventName, preventDefaults, false);
});

// Highlight drop zone when item is dragged over it
['dragenter', 'dragover'].forEach(eventName => {
    dropZone.addEventListener(eventName, () => dropZone.classList.add('highlight'), false);
});

['dragleave', 'drop'].forEach(eventName => {
    dropZone.addEventListener(eventName, () => dropZone.classList.remove('highlight'), false);
});

// Handle dropped files
dropZone.addEventListener('drop', handleDrop, false);

// Handle clear button click
const clearBtn = document.getElementById('clear-btn');
clearBtn.addEventListener('click', clearUploads);

// Handle file selection via click
fileUpload.addEventListener('change', handleFileSelect);

// Modal event listeners
showGraphBtn.addEventListener('click', () => {
    if (lastEmbeds) {
        graphModal.style.display = 'block';
        renderEmbeddingGraph(lastEmbeds.embeds1, lastEmbeds.embeds2);
    }
});

closeModalBtn.addEventListener('click', () => {
    graphModal.style.display = 'none';
});

window.addEventListener('click', (event) => {
    if (event.target == graphModal) {
        graphModal.style.display = 'none';
    }
});

function preventDefaults(e) {
    e.preventDefault();
    e.stopPropagation();
}

function handleDrop(e) {
    const dt = e.dataTransfer;
    const files = dt.files;
    handleFiles(files);
}

function handleFileSelect(e) {
    handleFiles(e.target.files);
}

function handleFiles(files) {
    const filesArray = Array.from(files);
    if (filesArray.length === 0) {
        return;
    }

    // If no image is uploaded yet, fill the first slot. Otherwise, fill the second.
    if (!imageSrc1) {
        handleIndividualFile(filesArray[0], '1');
        if (filesArray.length > 1) {
            handleIndividualFile(filesArray[1], '2');
        }
    } else {
        handleIndividualFile(filesArray[0], '2');
    }
}

function handleIndividualFile(file, target) {
    if (!file) {
        return;
    }
    
    const reader = new FileReader();
    reader.onload = function (e2) {
        const imageSrc = e2.target.result;

        if (target === '1') {
            imageSrc1 = imageSrc;
            imagePreview1.innerHTML = `<img src="${imageSrc1}" alt="uploaded image 1">`;
        } else if (target === '2') {
            imageSrc2 = imageSrc;
            imagePreview2.innerHTML = `<img src="${imageSrc2}" alt="uploaded image 2">`;
        }

        checkAndCompare();
    };
    reader.readAsDataURL(file);
}

function checkAndCompare() {
    if (imageSrc1 && imageSrc2) {
        compareImages(imageSrc1, imageSrc2);
    }
}

function clearUploads() {
    const placeholder = `<div class="placeholder">
        <i class="fas fa-image"></i>
        <p>Image preview</p>
    </div>`;
    imagePreview1.innerHTML = placeholder;
    imagePreview2.innerHTML = placeholder;

    imageSrc1 = null;
    imageSrc2 = null;

    lastEmbeds = null;
    showGraphBtn.style.display = 'none';
    // Reset file input
    fileUpload.value = '';

    meterContainer.innerHTML = '';
}

// Function to calculate cosine similarity between two vectors
function cosineSimilarity(vecA, vecB) {
    let dotProduct = 0.0;
    let normA = 0.0;
    let normB = 0.0;

    for (let i = 0; i < vecA.length; i++) {
        dotProduct += vecA[i] * vecB[i];
        normA += vecA[i] * vecA[i];
        normB += vecB[i] * vecB[i];
    }

    return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}

// Compare the two images using direct embedding calculation
async function compareImages(img1, img2) {
    statusText.textContent = 'Extracting embeddings...';
    spinner.style.display = 'block';

    try {
        // Load images using RawImage
        const image1 = await RawImage.read(img1);
        const image2 = await RawImage.read(img2);

        // Process images and compute embeddings
        const image_inputs1 = await processor(image1);
        const image_inputs2 = await processor(image2);

        const { image_embeds: embeds1 } = await vision_model(image_inputs1);
        const { image_embeds: embeds2 } = await vision_model(image_inputs2);

        // Calculate cosine similarity
        const similarity = cosineSimilarity(embeds1.data, embeds2.data);
        lastEmbeds = { embeds1: embeds1.data, embeds2: embeds2.data };

        statusText.textContent = 'Ready';
        spinner.style.display = 'none';
        renderResults(similarity);
    } catch (error) {
        statusText.textContent = '';
        spinner.style.display = 'none';
        meterContainer.innerHTML = `<div class="error"><p>Failed to compare images: ${error.message}</p></div>`;
        console.error('Comparison error:', error);
    }
}

// Render the comparison results
function renderResults(similarity) {
    meterContainer.innerHTML = '';

    // Show the button
    showGraphBtn.style.display = 'block';

    // Create similarity meter
    const meterElement = document.createElement('div'); // This will be a wrapper
    meterElement.className = 'similarity-meter';

    const score = Math.round(similarity * 100);
    const meterValue = Math.max(0, Math.min(100, score));

    meterElement.innerHTML = `
        <div class="meter-label">Similarity Score: ${meterValue}%</div>
        <div class="meter-container">
            <div class="meter-bar" style="width: ${meterValue}%"></div>
        </div>
        <div class="meter-description">${getSimilarityDescription(similarity)}</div>
    `;

    meterContainer.appendChild(meterElement);
}

function renderEmbeddingGraph(embeds1, embeds2) {
    graphContainerModal.innerHTML = `
        <h3 class="graph-title">Embedding Visualization</h3>
        <canvas id="embedding-chart"></canvas>
    `;

    const ctx = document.getElementById('embedding-chart').getContext('2d');
    new Chart(ctx, {
        type: 'line',
        data: {
            labels: Array.from({ length: embeds1.length }, (_, i) => i), // Dimension index
            datasets: [{
                label: 'Image 1 Embedding',
                data: embeds1,
                borderColor: 'rgba(0, 123, 255, 0.8)',
                backgroundColor: 'rgba(0, 123, 255, 0.1)',
                borderWidth: 1,
                pointRadius: 0,
            }, {
                label: 'Image 2 Embedding',
                data: embeds2,
                borderColor: 'rgba(40, 167, 69, 0.8)',
                backgroundColor: 'rgba(40, 167, 69, 0.1)',
                borderWidth: 1,
                pointRadius: 0,
            }]
        },
        options: {
            responsive: true,
            plugins: {
                legend: { position: 'top' },
                zoom: {
                    pan: {
                        enabled: true,
                        mode: 'xy',
                        modifierKey: null, // Allow panning without holding a key
                    },
                    zoom: {
                        wheel: {
                            enabled: true,
                        },
                        pinch: {
                            enabled: true
                        },
                        mode: 'xy',
                    }
                }
            }
        }
    });
}

function getSimilarityDescription(similarity) {
    if (similarity > 0.9) {
        return "🔥 Extremely similar - These images are nearly identical!";
    } else if (similarity > 0.7) {
        return "👍 Very similar - These images share strong visual characteristics.";
    } else {
        return "🚫 Not similar - These images appear to be very different.";
    }
}