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

// DOM Elements
const chatMessages = document.getElementById('chatMessages');
const userInput = document.getElementById('userInput');
const sendButton = document.getElementById('sendButton');
const modelStatus = document.getElementById('modelStatus');
const statusText = document.getElementById('statusText');
const progressFill = document.getElementById('progressFill');

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

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

// Create text generation pipeline with progress tracking
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 percentage = Math.round((progress.loaded / progress.total) * 100);
updateStatus(`Loading model: ${progress.file}`, percentage);
}
}
}
);

updateStatus('Model ready!', 100);

// Enable input after model is loaded
setTimeout(() => {
modelStatus.classList.add('ready');
userInput.disabled = false;
sendButton.disabled = false;
userInput.focus();
}, 500);

} catch (error) {
console.error('Error initializing model:', error);
updateStatus('Error loading model. Please refresh the page.', 0);
statusText.style.color = '#FF3B30';
}
}

// Update status display
function updateStatus(message, progress) {
statusText.textContent = message;
progressFill.style.width = `${progress}%`;
}

// Add message to chat
function addMessage(role, content, isStreaming = false) {
const messageDiv = document.createElement('div');
messageDiv.className = `message ${role}`;

const avatar = document.createElement('div');
avatar.className = 'message-avatar';
avatar.textContent = role === 'user' ? '👤' : '🤖';

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

if (isStreaming) {
contentDiv.innerHTML = '<div class="typing-indicator">
  <div class="typing-dot"></div>
  <div class="typing-dot"></div>
  <div class="typing-dot"></div>
</div>';
} else {
contentDiv.innerHTML = formatMessage(content);
}

messageDiv.appendChild(avatar);
messageDiv.appendChild(contentDiv);

// Remove welcome message if it exists
const welcomeMessage = chatMessages.querySelector('.welcome-message');
if (welcomeMessage) {
welcomeMessage.remove();
}

chatMessages.appendChild(messageDiv);
chatMessages.scrollTop = chatMessages.scrollHeight;

return contentDiv;
}

// Format message content
function formatMessage(text) {
// Simple formatting for line breaks
return text
.split('\n')
.map(line => line.trim() ? `<p>${escapeHtml(line)}</p>` : '')
.join('');
}

// Escape HTML to prevent XSS
function escapeHtml(text) {
const div = document.createElement('div');
div.textContent = text;
return div.innerHTML;
}

// Handle message sending
async function sendMessage() {
const message = userInput.value.trim();
if (!message || isGenerating || !generator) return;

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

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

// Disable input while generating
isGenerating = true;
userInput.disabled = true;
sendButton.disabled = true;

try {
// Add assistant message placeholder with typing indicator
const assistantContentDiv = addMessage('assistant', '', true);
let fullResponse = '';

// Create streamer for real-time output
const streamer = new TextStreamer(generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
fullResponse += text;
assistantContentDiv.innerHTML = formatMessage(fullResponse);
chatMessages.scrollTop = chatMessages.scrollHeight;
}
});

// Generate response
const output = await generator(conversationHistory, {
max_new_tokens: 512,
do_sample: true,
temperature: 0.7,
top_p: 0.9,
streamer: streamer
});

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

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

// Ensure final response is displayed
assistantContentDiv.innerHTML = formatMessage(finalResponse);

} catch (error) {
console.error('Error generating response:', error);
addMessage('assistant', 'Sorry, I encountered an error. Please try again.');
} finally {
// Re-enable input
isGenerating = false;
userInput.disabled = false;
sendButton.disabled = false;
userInput.focus();
}
}

// Auto-resize textarea
userInput.addEventListener('input', function() {
this.style.height = 'auto';
this.style.height = Math.min(this.scrollHeight, 120) + 'px';
});

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

// Send button click
sendButton.addEventListener('click', sendMessage);

// Initialize app when page loads
init();