import { pipeline, TextStreamer } from '@huggingface/transformers';

// DOM Elements
const chatHistory = document.getElementById('chat-history');
const userInput = document.getElementById('user-input');
const sendButton = document.getElementById('send-button');

// Global pipeline variable
let generator = null;

// Initialize the model when the page loads
async function initializeModel() {
  try {
    showLoader('Loading model...');
    generator = await pipeline('text-generation', 'onnx-community/Qwen3-0.6B-ONNX', { dtype: 'q4f16' });
    hideLoader();
    addMessage('System', 'Model loaded successfully! You can now start chatting.', 'system');
  } catch (error) {
    hideLoader();
    addMessage('System', `Error loading model: ${error.message}`, 'error');
  }
}

// Show loading message
function showLoader(message) {
  const loaderElement = document.createElement('div');
  loaderElement.className = 'chat-message system';
  loaderElement.id = 'loader';
  loaderElement.innerHTML = `
    <div class="message-content">
      <div class="loader"></div>
      <span>${message}</span>
    </div>
  `;
  chatHistory.appendChild(loaderElement);
  chatHistory.scrollTop = chatHistory.scrollHeight;
}

// Hide loading message
function hideLoader() {
  const loaderElement = document.getElementById('loader');
  if (loaderElement) loaderElement.remove();
}

// Add message to chat display
function addMessage(role, content, type = 'user') {
  const messageElement = document.createElement('div');
  messageElement.className = `chat-message ${type}`;
  
  if (type === 'user' || type === 'assistant') {
    messageElement.innerHTML = `<strong>${role}:</strong> ${content}`;
  } else {
    messageElement.innerHTML = `<em>${content}</em>`;
  }
  
  chatHistory.appendChild(messageElement);
  chatHistory.scrollTop = chatHistory.scrollHeight;
}

// Handle sending messages
async function handleSendMessage() {
  if (!generator) {
    addMessage('System', 'Model is still loading. Please wait...', 'error');
    return;
  }

  const input = userInput.value.trim();
  if (!input) return;

  // Add user message to chat
  addMessage('You', input, 'user');
  userInput.value = '';

  try {
    // Prepare messages for model
    const messages = [
      { role: 'system', content: 'You are a helpful assistant.' },
      { role: 'user', content: input }
    ];

    // Create temporary element for streaming response
    const assistantMessageElement = document.createElement('div');
    assistantMessageElement.className = 'chat-message assistant';
    assistantMessageElement.innerHTML = '<strong>Assistant:</strong> ';
    const contentSpan = document.createElement('span');
    assistantMessageElement.appendChild(contentSpan);
    chatHistory.appendChild(assistantMessageElement);

    // Create streamer to handle response tokens
    const streamer = new TextStreamer(generator.tokenizer, {
      skip_prompt: true,
      skip_special_tokens: true,
      callback_function: (token) => {
        contentSpan.textContent += token;
        chatHistory.scrollTop = chatHistory.scrollHeight;
      }
    });

    // Generate response
    await generator(messages, {
      max_new_tokens: 256,
      do_sample: false,
      streamer: streamer
    });

    // Final response element update
    contentSpan.textContent = contentSpan.textContent.trim();
    chatHistory.scrollTop = chatHistory.scrollHeight;
  } catch (error) {
    addMessage('System', `Error generating response: ${error.message}`, 'error');
  }
}

// Event Listeners
sendButton.addEventListener('click', handleSendMessage);

userInput.addEventListener('keydown', (e) => {
  if (e.key === 'Enter' && !e.shiftKey) {
    e.preventDefault();
    handleSendMessage();
  }
});

// Initialize the model on page load
document.addEventListener('DOMContentLoaded', initializeModel);