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

document.addEventListener('DOMContentLoaded', () => {
    const chatContainer = document.getElementById('chat-container');
    const chatForm = document.getElementById('chat-form');
    const userInput = document.getElementById('user-input');
    const loading = document.getElementById('loading');
    const errorDiv = document.getElementById('error');
    let generator = null;
    let messages = [{ role: "system", content: "You are a helpful assistant." }];

    async function initGenerator() {
        try {
            loading.classList.remove('hidden');
            errorDiv.classList.add('hidden');
            generator = await pipeline(
                "text-generation",
                "onnx-community/gemma-3-270m-it-ONNX",
                { dtype: "fp32" }
            );
            loading.classList.add('hidden');
            addMessage('system', 'Chatbot is ready! How can I help you?');
        } catch (err) {
            loading.classList.add('hidden');
            errorDiv.textContent = 'Failed to load the model. Please try refreshing the page.';
            errorDiv.classList.remove('hidden');
            console.error(err);
        }
    }

    function addMessage(role, content) {
        const messageDiv = document.createElement('div');
        messageDiv.classList.add('message', role === 'user' ? 'user-message' : 'bot-message');
        messageDiv.textContent = content;
        chatContainer.appendChild(messageDiv);
        chatContainer.scrollTop = chatContainer.scrollHeight;
    }

    async function generateResponse() {
        if (!generator) return;

        loading.classList.remove('hidden');
        document.getElementById('send-btn').disabled = true;
        userInput.disabled = true;

        const botMessageDiv = document.createElement('div');
        botMessageDiv.classList.add('message', 'bot-message');
        chatContainer.appendChild(botMessageDiv);
        chatContainer.scrollTop = chatContainer.scrollHeight;

        try {
            const streamer = new TextStreamer(generator.tokenizer, {
                skip_prompt: true,
                skip_special_tokens: true,
                callback_function: (text) => {
                    botMessageDiv.textContent += text;
                    chatContainer.scrollTop = chatContainer.scrollHeight;
                }
            });

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

            const response = output[0].generated_text.at(-1).content;
            messages.push({ role: "assistant", content: response });
        } catch (err) {
            botMessageDiv.textContent = 'Error generating response.';
            console.error(err);
        } finally {
            loading.classList.add('hidden');
            document.getElementById('send-btn').disabled = false;
            userInput.disabled = false;
            userInput.focus();
        }
    }

    chatForm.addEventListener('submit', (e) => {
        e.preventDefault();
        const userText = userInput.value.trim();
        if (!userText) return;

        addMessage('user', userText);
        messages.push({ role: "user", content: userText });
        userInput.value = '';
        generateResponse();
    });

    initGenerator();
});