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

// DOM Elements
const messagesContainer = document.getElementById('messages');
const userInput = document.getElementById('userInput');
const sendBtn = document.getElementById('sendBtn');
const status = document.getElementById('status');
const loadingOverlay = document.getElementById('loadingOverlay');
const loadingText = document.getElementById('loadingText');
const progressFill = document.getElementById('progressFill');
const progressText = document.getElementById('progressText');

// State
let generator = null;
let conversationHistory = [
{ role: "system", content: "You are a helpful, friendly, and knowledgeable AI assistant. Provide clear, concise, and
accurate responses." }
];
let isGenerating = false;

// Initialize the application
async function initializeApp() {
try {
updateLoadingStatus('Loading AI model...', 0);

// Create a text generation pipeline with progress callback
generator = await pipeline(
"text-generation",
"onnx-community/Llama-3.2-1B-Instruct-q4f16",
{
dtype: "q4f16",
device: "webgpu",
progress_callback: (progress) => {
if (progress.status === 'progress') {
const percent = Math.round((progress.loaded / progress.total) * 100);
updateLoadingStatus(`Loading ${progress.file}...`, percent);
} else if (progress.status === 'done') {
updateLoadingStatus('Initializing model...', 95);
}
}
}
);

updateLoadingStatus('Ready!', 100);

// Hide loading overlay after a short delay
setTimeout(() => {
loadingOverlay.classList.add('hidden');
enableChat();
}, 500);

} catch (error) {
console.error('Error initializing model:', error);
updateLoadingStatus('Error loading model. Please refresh the page.', 0);
status.textContent = 'Failed to load AI model. Please refresh.';
status.style.color = '#FF3B30';
}
}

function updateLoadingStatus(text, percent) {
loadingText.textContent = text;
progressFill.style.width = `${percent}%`;
progressText.textContent = `${percent}%`;
}

function enableChat() {
userInput.disabled = false;
sendBtn.disabled = false;
status.textContent = 'Ready to chat';
userInput.focus();
}

function addMessage(role, content) {
// Remove welcome message if it exists
const welcomeMessage = messagesContainer.querySelector('.welcome-message');
if (welcomeMessage) {
welcomeMessage.remove();
}

const messageDiv = document.createElement('div');
messageDiv.className = `message ${role}`;

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

messageDiv.appendChild(contentDiv);
messagesContainer.appendChild(messageDiv);

// Scroll to bottom
messagesContainer.scrollTop = messagesContainer.scrollHeight;

return contentDiv;
}

function addTypingIndicator() {
const messageDiv = document.createElement('div');
messageDiv.className = 'message assistant';
messageDiv.id = 'typing-indicator';

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

const typingDiv = document.createElement('div');
typingDiv.className = 'typing-indicator';
typingDiv.innerHTML = `
<div class="typing-dot"></div>
<div class="typing-dot"></div>
<div class="typing-dot"></div>
`;

contentDiv.appendChild(typingDiv);
messageDiv.appendChild(contentDiv);
messagesContainer.appendChild(messageDiv);

messagesContainer.scrollTop = messagesContainer.scrollHeight;

return messageDiv;
}

function removeTypingIndicator() {
const indicator = document.getElementById('typing-indicator');
if (indicator) {
indicator.remove();
}
}

async function generateResponse(userMessage) {
if (isGenerating) return;

isGenerating = true;
sendBtn.disabled = true;
userInput.disabled = true;
status.textContent = 'Thinking...';

// Add user message to conversation history
conversationHistory.push({ role: "user", content: userMessage });

// Show typing indicator
const typingIndicator = addTypingIndicator();

try {
let assistantMessage = '';
let messageElement = null;

// Create a custom streamer with callback
const streamer = new TextStreamer(generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
// Remove typing indicator on first token
if (!messageElement) {
removeTypingIndicator();
messageElement = addMessage('assistant', '');
}

// Append the new text
assistantMessage += text;
messageElement.textContent = assistantMessage;

// Scroll to bottom
messagesContainer.scrollTop = messagesContainer.scrollHeight;
}
});

// Generate response with streaming
const output = await generator(conversationHistory, {
max_new_tokens: 512,
do_sample: false,
temperature: 0.7,
streamer: streamer,
});

// Get the final response
const finalResponse = output[0].generated_text.at(-1).content;

// Update conversation history
conversationHistory.push({ role: "assistant", content: finalResponse });

// Ensure the final message is displayed
if (messageElement) {
messageElement.textContent = finalResponse;
}

status.textContent = 'Ready to chat';

} catch (error) {
console.error('Error generating response:', error);
removeTypingIndicator();
addMessage('assistant', 'Sorry, I encountered an error. Please try again.');
status.textContent = 'Error occurred';
} finally {
isGenerating = false;
sendBtn.disabled = false;
userInput.disabled = false;
userInput.focus();
}
}

function handleSend() {
const message = userInput.value.trim();

if (!message || isGenerating) return;

// Add user message to UI
addMessage('user', message);

// Clear input
userInput.value = '';
userInput.style.height = 'auto';

// Generate response
generateResponse(message);
}

// Event Listeners
sendBtn.addEventListener('click', handleSend);

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

// Auto-resize textarea
userInput.addEventListener('input', () => {
userInput.style.height = 'auto';
userInput.style.height = userInput.scrollHeight + 'px';
});

// Initialize the app
initializeApp();