import { AutoProcessor, AutoTokenizer, Moondream1ForConditionalGeneration, RawImage } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@2.17.2';

class MoondreamAssistant {
    constructor() {
        this.model = null;
        this.processor = null;
        this.tokenizer = null;
        this.currentImage = null;
        this.isModelLoaded = false;
        
        this.initializeElements();
        this.bindEvents();
        this.loadModel();
    }

    initializeElements() {
        this.uploadArea = document.getElementById('uploadArea');
        this.fileInput = document.getElementById('fileInput');
        this.imagePreview = document.getElementById('imagePreview');
        this.previewImage = document.getElementById('previewImage');
        this.removeBtn = document.getElementById('removeBtn');
        this.querySection = document.getElementById('querySection');
        this.questionInput = document.getElementById('questionInput');
        this.askBtn = document.getElementById('askBtn');
        this.responseSection = document.getElementById('responseSection');
        this.loading = document.getElementById('loading');
        this.responseText = document.getElementById('responseText');
    }

    bindEvents() {
        this.uploadArea.addEventListener('click', () => this.fileInput.click());
        this.uploadArea.addEventListener('dragover', this.handleDragOver.bind(this));
        this.uploadArea.addEventListener('drop', this.handleDrop.bind(this));
        this.fileInput.addEventListener('change', this.handleFileSelect.bind(this));
        this.removeBtn.addEventListener('click', this.removeImage.bind(this));
        this.questionInput.addEventListener('input', this.handleQuestionInput.bind(this));
        this.askBtn.addEventListener('click', this.askQuestion.bind(this));
        this.questionInput.addEventListener('keypress', (e) => {
            if (e.key === 'Enter' && !this.askBtn.disabled) {
                this.askQuestion();
            }
        });
    }

    async loadModel() {
        try {
            const model_id = 'Xenova/moondream2';
            
            this.updateUploadArea('Loading AI model...', true);
            
            [this.processor, this.tokenizer, this.model] = await Promise.all([
                AutoProcessor.from_pretrained(model_id),
                AutoTokenizer.from_pretrained(model_id),
                Moondream1ForConditionalGeneration.from_pretrained(model_id, {
                    dtype: {
                        embed_tokens: 'fp16',
                        vision_encoder: 'fp16',
                        decoder_model_merged: 'q4',
                    },
                    device: 'webgpu',
                })
            ]);

            this.isModelLoaded = true;
            this.updateUploadArea('Drop an image here or browse');
            console.log('Model loaded successfully');
        } catch (error) {
            console.error('Error loading model:', error);
            this.updateUploadArea('Error loading AI model. Please refresh.', true);
        }
    }

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

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

    handleFileSelect(e) {
        const file = e.target.files[0];
        if (file) {
            this.handleImageFile(file);
        }
    }

    async handleImageFile(file) {
        if (!file.type.startsWith('image/')) {
            alert('Please select a valid image file');
            return;
        }

        try {
            this.currentImage = await RawImage.fromBlob(file);
            const url = URL.createObjectURL(file);
            
            this.previewImage.src = url;
            this.imagePreview.style.display = 'block';
            this.uploadArea.style.display = 'none';
            this.querySection.style.display = 'block';
            this.questionInput.focus();
        } catch (error) {
            console.error('Error loading image:', error);
            alert('Error loading image. Please try again.');
        }
    }

    removeImage() {
        this.currentImage = null;
        this.imagePreview.style.display = 'none';
        this.uploadArea.style.display = 'flex';
        this.querySection.style.display = 'none';
        this.responseSection.style.display = 'none';
        this.questionInput.value = '';
        this.askBtn.disabled = true;
        
        if (this.previewImage.src) {
            URL.revokeObjectURL(this.previewImage.src);
        }
    }

    handleQuestionInput() {
        this.askBtn.disabled = this.questionInput.value.trim().length === 0;
    }

    async askQuestion() {
        if (!this.isModelLoaded || !this.currentImage) return;

        const question = this.questionInput.value.trim();
        if (!question) return;

        this.askBtn.disabled = true;
        this.loading.style.display = 'flex';
        this.responseSection.style.display = 'block';
        this.responseText.textContent = '';

        try {
            const text = `<image>\n\nQuestion: ${question}\n\nAnswer:`;
            const text_inputs = this.tokenizer(text);
            const vision_inputs = await this.processor(this.currentImage);

            const output = await this.model.generate({
                ...text_inputs,
                ...vision_inputs,
                do_sample: false,
                max_new_tokens: 128,
            });

            const decoded = this.tokenizer.batch_decode(output, { skip_special_tokens: true });
            const response = decoded[0]
                .split('Answer:')[1]
                .trim()
                .replace(/<\|.*?\|>/g, '');

            this.responseText.textContent = response;
        } catch (error) {
            console.error('Error generating response:', error);
            this.responseText.textContent = 'Sorry, there was an error processing your question. Please try again.';
        } finally {
            this.loading.style.display = 'none';
            this.askBtn.disabled = false;
        }
    }

    updateUploadArea(text, isError = false) {
        const p = this.uploadArea.querySelector('p');
        p.textContent = text;
        this.uploadArea.style.pointerEvents = isError ? 'none' : 'auto';
    }
}

// Initialize the application
document.addEventListener('DOMContentLoaded', () => {
    new MoondreamAssistant();
});