import { pipeline, TextStreamer } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.1';

class GemmaChatbot {
    constructor() {
        this.generator = null;
        this.messages = [
            { role: "system", content: "You are a helpful, friendly AI assistant. Keep your responses concise and helpful." }
        ];
        this.isGenerating = false;
        this.initializeElements();
        this.initializeModel();
    }

    initializeElements() {
        this.chatMessages = document.getElementById('chatMessages');
        this.userInput = document.getElementById('userInput');
        this.sendButton = document.getElementById('sendButton');
        this.loadingIndicator = document.getElementById('loadingIndicator');

        this.sendButton.addEventListener('click', () => this.sendMessage());
        this.userInput.addEventListener('keypress', (e) => {
            if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                this.sendMessage();
            }
        });
    }

    async initializeModel() {
        try {
            console.log('Initializing Gemma model...');
            this.generator = await pipeline(
                'text-generation',
                'onnx-community/gemma-3-270m-it-ONNX',
                { dtype: 'fp32' }
            );
            console.log('Model loaded successfully');
            this.enableChat();
        } catch (error) {
            console.error('Failed to initialize model:', error);
            this.showError('Failed to load AI model. Please refresh the page to try again.');
        }
    }

    enableChat() {
        this.loadingIndicator.style.display = 'none';
        this.userInput.disabled = false;
        this.sendButton.disabled = false;
        this.userInput.focus();
    }

    addMessage(content, isUser = false) {
        const messageDiv = document.createElement('div');
        messageDiv.className = `message ${isUser ? 'user' : 'assistant'}`;
        
        const contentDiv = document.createElement('div');
        contentDiv.className = 'message-content';
        contentDiv.textContent = content;
        
        messageDiv.appendChild(contentDiv);
        this.chatMessages.appendChild(messageDiv);
        this.scrollToBottom();
        
        return contentDiv;
    }

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

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

        this.isGenerating = true;
        this.userInput.value = '';
        this.userInput.disabled = true;
        this.sendButton.disabled = true;

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

        // Add assistant message placeholder
        const assistantMessageDiv = document.createElement('div');
        assistantMessageDiv.className = 'message assistant';
        
        const contentDiv = document.createElement('div');
        contentDiv.className = 'message-content';
        
        const typingIndicator = document.createElement('span');
        typingIndicator.className = 'typing-indicator';
        typingIndicator.innerHTML = '<span></span><span></span><span></span>';
        contentDiv.appendChild(typingIndicator);
        
        assistantMessageDiv.appendChild(contentDiv);
        this.chatMessages.appendChild(assistantMessageDiv);
        this.scrollToBottom();

        try {
            let fullResponse = '';
            
            // Create custom streamer
            const streamer = new TextStreamer(this.generator.tokenizer, {
                skip_prompt: true,
                skip_special_tokens: true,
                callback_function: (text) => {
                    fullResponse += text;
                    contentDiv.textContent = fullResponse;
                    this.scrollToBottom();
                }
            });

            // Generate response
            const output = await this.generator(this.messages, {
                max_new_tokens: 256,
                do_sample: true,
                temperature: 0.7,
                top_p: 0.9,
                streamer: streamer
            });

            const generatedContent = output[0].generated_text.at(-1).content;
            this.messages.push({ role: "assistant", content: generatedContent });
            
            // Ensure final content is displayed
            contentDiv.textContent = generatedContent;
            
        } catch (error) {
            console.error('Generation error:', error);
            contentDiv.textContent = 'Sorry, I encountered an error. Please try again.';
        } finally {
            this.isGenerating = false;
            this.userInput.disabled = false;
            this.sendButton.disabled = false;
            this.userInput.focus();
        }
    }

    showError(message) {
        this.loadingIndicator.innerHTML = `
            <div class="error-message">
                <p>⚠️ ${message}</p>
            </div>
        `;
    }
}

// Initialize chatbot when DOM is ready
document.addEventListener('DOMContentLoaded', () => {
    new GemmaChatbot();
});