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

        this.initElements();
        this.bindEvents();
        this.loadModel();
    }

    initElements() {
        this.chatContainer = document.getElementById('chat-container');
        this.messageInput = document.getElementById('message-input');
        this.sendButton = document.getElementById('send-button');
        this.loadingIndicator = document.getElementById('loading-indicator');
        this.deviceSelect = document.getElementById('device-select');
        this.errorToast = document.getElementById('error-toast');
        this.errorMessage = document.getElementById('error-message');
    }

    bindEvents() {
        this.sendButton.addEventListener('click', () => this.sendMessage());
        this.messageInput.addEventListener('keypress', (e) => {
            if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                this.sendMessage();
            }
        });
        this.messageInput.addEventListener('input', this.autoResize.bind(this));
        this.deviceSelect.addEventListener('change', (e) => {
            this.device = e.target.value;
            this.loadModel();
        });
    }

    async loadModel() {
        try {
            this.showLoading();
            this.sendButton.disabled = true;
            this.messageInput.disabled = true;

            const config = {
                dtype: "fp32"
            };

            if (this.device === 'webgpu') {
                config.device = 'webgpu';
            }

            this.generator = await window.transformers.pipeline(
                "text-generation",
                "onnx-community/gemma-3-270m-it-ONNX",
                config
            );

            this.hideLoading();
            this.showMessage('system', 'Model loaded successfully! Ready to chat.', 'info');
        } catch (error) {
            console.error('Failed to load model:', error);
            this.showError(`Failed to load model: ${error.message}. Try CPU mode or refresh the page.`);
            if (this.device === 'webgpu') {
                this.deviceSelect.value = 'cpu';
                this.device = 'cpu';
            }
        } finally {
            this.sendButton.disabled = false;
            this.messageInput.disabled = false;
        }
    }

    async sendMessage() {
        const message = this.messageInput.value.trim();
        if (!message || this.isGenerating || !this.generator) return;

        this.messageInput.value = '';
        this.autoResize();

        // Add user message
        this.conversation.push({ role: "user", content: message });
        this.addMessage('user', message);

        this.isGenerating = true;
        this.showLoading();
        this.sendButton.disabled = true;
        this.messageInput.disabled = true;

        try {
            const streamer = new window.transformers.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).content;
            this.conversation.push({ role: "assistant", content: assistantMessage });
            this.addMessage('assistant', assistantMessage);
        } catch (error) {
            console.error('Generation error:', error);
            this.showError('Sorry, an error occurred while generating the response.');
            this.addMessage('assistant', "Sorry, I encountered an error. Please try again.");
        } finally {
            this.isGenerating = false;
            this.hideLoading();
            this.sendButton.disabled = false;
            this.messageInput.disabled = false;
            this.messageInput.focus();
        }
    }

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

        const messageContent = document.createElement('div');
        messageContent.className = 'message-content';

        const roleName = role === 'user' ? 'You' : 'Gemma';
        messageContent.innerHTML = `<strong>${roleName}:</strong> ${this.escapeHtml(content)}`;

        messageDiv.appendChild(messageContent);
        this.chatContainer.appendChild(messageDiv);
        this.scrollToBottom();
    }

    showMessage(role, content, type) {
        this.addMessage(role, content, type);
    }

    autoResize() {
        this.messageInput.style.height = 'auto';
        this.messageInput.style.height = Math.min(this.messageInput.scrollHeight, 120) + 'px';
    }

    scrollToBottom() {
        this.chatContainer.scrollTop = this.chatContainer.scrollHeight;
    }

    showLoading() {
        this.loadingIndicator.classList.remove('hidden');
    }

    hideLoading() {
        this.loadingIndicator.classList.add('hidden');
    }

    showError(message) {
        this.errorMessage.textContent = message;
        this.errorToast.classList.remove('hidden');
        setTimeout(() => {
            this.errorToast.classList.add('hidden');
        }, 5000);
    }

    escapeHtml(text) {
        const div = document.createElement('div');
        div.textContent = text;
        return div.innerHTML;
    }
}

// Initialize the chatbot when the page loads
document.addEventListener('DOMContentLoaded', () => {
    new GemmaChatbot();
});