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

// Global variables
let generator = null;
let conversationHistory = [];
let isGenerating = false;

// DOM elements
const messagesContainer = document.getElementById('messagesContainer');
const userInput = document.getElementById('userInput');
const sendButton = document.getElementById('sendButton');
const clearButton = document.getElementById('clearButton');
const statusText = document.getElementById('statusText');
const progressContainer = document.getElementById('progressContainer');
const progressBar = document.getElementById('progressBar');

// Initialize the application
async function initialize() {
try {
updateStatus('Loading model... This may take 1-2 minutes on first load.', true);

// Create a 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 === 'downloading') {
const percent = Math.round((progress.loaded / progress.total) * 100);
updateProgress(percent);
updateStatus(`Downloading model: ${progress.file} (${percent}%)`, true);
} else if (progress.status === 'loading') {
updateStatus(`Loading model into memory...`, true);
}
}
}
);

updateStatus('✓ Model loaded successfully! Ready to chat.', false);
enableInput();

// Initialize conversation with system message
conversationHistory = [
{ role: "system", content: "You are a helpful, friendly, and knowledgeable AI assistant. Provide clear, concise, and
accurate responses." }
];

} catch (error) {
console.error('Initialization error:', error);
updateStatus(`❌ Error loading model: ${error.message}`, false);
showError('Failed to load the model. Please refresh the page and try again.');
}
}

// Update status bar
function updateStatus(message, showProgress) {
statusText.textContent = message;
progressContainer.style.display = showProgress ? 'block' : 'none';
}

// Update progress bar
function updateProgress(percent) {
progressBar.style.width = `${percent}%`;
}

// Enable input controls
function enableInput() {
userInput.disabled = false;
sendButton.disabled = false;
userInput.focus();
}

// Disable input controls
function disableInput() {
userInput.disabled = true;
sendButton.disabled = true;
}

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

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

const roleLabel = document.createElement('strong');
roleLabel.textContent = role === 'user' ? 'You:' : 'Assistant:';

const messageText = document.createElement('p');
messageText.textContent = content;

contentDiv.appendChild(roleLabel);
contentDiv.appendChild(messageText);
messageDiv.appendChild(contentDiv);

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

return messageText;
}

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

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

const indicator = document.createElement('div');
indicator.className = 'typing-indicator';
indicator.innerHTML = '<span></span><span></span><span></span>';

contentDiv.appendChild(indicator);
messageDiv.appendChild(contentDiv);
messagesContainer.appendChild(messageDiv);
messagesContainer.scrollTop = messagesContainer.scrollHeight;
}

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

// Generate response
async function generateResponse(userMessage) {
if (!generator || isGenerating) return;

isGenerating = true;
disableInput();

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

// Show typing indicator
addTypingIndicator();

try {
let assistantResponse = '';
const messageElement = document.createElement('p');

// Create streamer with callback
const streamer = new TextStreamer(generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
assistantResponse += text;
removeTypingIndicator();

// Update or create message element
if (!messageElement.parentElement) {
const messageDiv = document.createElement('div');
messageDiv.className = 'message assistant';

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

const roleLabel = document.createElement('strong');
roleLabel.textContent = 'Assistant:';

contentDiv.appendChild(roleLabel);
contentDiv.appendChild(messageElement);
messageDiv.appendChild(contentDiv);
messagesContainer.appendChild(messageDiv);
}

messageElement.textContent = assistantResponse;
messagesContainer.scrollTop = messagesContainer.scrollHeight;
}
});

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

// Add assistant response to history
const finalResponse = output[0].generated_text.at(-1).content;
conversationHistory.push({ role: "assistant", content: finalResponse });

updateStatus('✓ Model loaded successfully! Ready to chat.', false);

} catch (error) {
console.error('Generation error:', error);
removeTypingIndicator();
addMessage('assistant', `Sorry, I encountered an error: ${error.message}`);
updateStatus('❌ Error generating response', false);
} finally {
isGenerating = false;
enableInput();
}
}

// Handle send button click
async function handleSend() {
const message = userInput.value.trim();
if (!message || isGenerating) return;

userInput.value = '';
userInput.style.height = 'auto';
await generateResponse(message);
}

// Handle clear button click
function handleClear() {
if (confirm('Are you sure you want to clear the conversation?')) {
// Keep only system message
conversationHistory = conversationHistory.slice(0, 1);

// Clear UI
messagesContainer.innerHTML = '';
addMessage('assistant', "Hello! I'm your AI assistant powered by Gemma. How can I help you today?");

userInput.value = '';
userInput.focus();
}
}

// Show error message
function showError(message) {
const errorDiv = document.createElement('div');
errorDiv.className = 'message assistant';
errorDiv.innerHTML = `
<div class="message-content" style="background: #fee; color: #c33;">
  <strong>Error:</strong>
  <p>${message}</p>
</div>
`;
messagesContainer.appendChild(errorDiv);
}

// Event listeners
sendButton.addEventListener('click', handleSend);
clearButton.addEventListener('click', handleClear);

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

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

// Initialize on load
initialize();