// index.js
class GemmaChatbot {
    constructor() {
        this.generator = null;
        this.conversation = [
            { role: "system", content: "You are a helpful assistant." }
        ];
        this.isGenerating = false;
        this.useWebGPU = false;
        
        this.initializeElements();
        this.bindEvents();
        this.checkWebGPU();
    }

    initializeElements() {
        this.chatContainer = document.getElementById('chatContainer');
        this.messageInput = document.getElementById('messageInput');
        this.sendButton = document.getElementById('sendButton');
        this.deviceToggle = document.getElementById('device-toggle');
        this.statusIndicator = document.getElementById('statusIndicator');
        this.statusText = document.querySelector('.status-text');
        this.statusDots = document.getElementById('statusDots');
    }

    bindEvents() {
        this.messageInput.addEventListener('input', () => this.handleInput());
        this.messageInput.addEventListener('keydown', (e) => this.handleKeydown(e));
        this.sendButton.addEventListener('click', () => this.sendMessage());
        this.deviceToggle.addEventListener('change', (e) => this.toggleDevice(e.target.checked));
        
        // Auto-resize textarea
        this.messageInput.addEventListener('input', () => {
            this.messageInput.style.height = 'auto';
            this.messageInput.style.height = Math.min(this.messageInput.scrollHeight, 120) + 'px';
        });
    }

    async initModel() {
        try {
            this.setStatus('Loading model...', true);
            
            const options = { dtype: "fp32" };
            if (this.useWebGPU) {
                options.device = 'webgpu';
            }
            
            this.generator = await window.pipeline(
                "text-generation",
                "onnx-community/gemma-3-270m-it-ONNX",
                options
            );
            
            this.setStatus('Ready', false);
            console.log('Model loaded successfully');
        } catch (error) {
            console.error('Failed to load model:', error);
            this.setStatus('Failed to load model. Retrying...', false);
            setTimeout(() => this.initModel(), 3000);
        }
    }

    async sendMessage() {
        if (this.isGenerating || !this.generator) return;

        const message = this.messageInput.value.trim();
        if (!message) return;

        this.addMessage('user', message);
        this.messageInput.value = '';
        this.messageInput.style.height = 'auto';
        this.setInputEnabled(false);

        this.conversation.push({ role: "user", content: message });
        this.isGenerating = true;
        this.setStatus('Generating response...', true);

        try {
            const streamer = new window.TextStreamer(this.generator.tokenizer, {
                skip_prompt: true,
                skip_special_tokens: true,
            });

            const output = await this.generator(this.conversation, {
                max_new_tokens: 512,
                do_sample: false,
                streamer: streamer,
            });

            const assistantMessage = output[0].generated_text.at(-1);
            this.conversation.push(assistantMessage);
            this.addMessage('assistant', assistantMessage.content);
            
        } catch (error) {
            console.error('Generation error:', error);
            this.addMessage('assistant', 'Sorry, I encountered an error. Please try again.');
        } finally {
            this.isGenerating = false;
            this.setInputEnabled(true);
            this.setStatus('Ready', false);
        }
    }

    addMessage(role, content) {
        const messageDiv = document.createElement('div');
        messageDiv.className = `message-wrapper ${role}-message-wrapper`;

        const avatar = document.createElement('div');
        avatar.className = `avatar ${role}-avatar`;
        avatar.textContent = role