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

// Configure environment
env.allowLocalModels = false;

class ChatApp {
constructor() {
this.generator = null;
this.conversationHistory = [
{ role: "system", content: "You are a helpful, friendly AI assistant. Provide clear, concise, and accurate responses." }
];
this.isGenerating = false;

this.elements = {
loadingScreen: document.getElementById('loadingScreen'),
chatInterface: document.getElementById('chatInterface'),
messagesContainer: document.getElementById('messagesContainer'),
userInput: document.getElementById('userInput'),
sendButton: document.getElementById('sendButton'),
loadingStatus: document.getElementById('loadingStatus'),
progressFill: document.getElementById('progressFill'),
progressText: document.getElementById('progressText')
};

this.init();
}

async init() {
try {
await this.loadModel();
this.setupEventListeners();
this.showChatInterface();
} catch (error) {
console.error('Initialization error:', error);
this.showError('Failed to initialize the application. Please refresh the page.');
}
}

async loadModel() {
this.updateLoadingStatus('Downloading model...', 0);

try {
this.generator = await pipeline(
"text-generation",
"onnx-community/Qwen2.5-0.5B-Instruct",
{
dtype: "q8",
device: "webgpu",
progress_callback: (progress) => {
if (progress.status === 'progress') {
const percent = Math.round((progress.loaded / progress.total) * 100);
this.updateLoadingStatus(
`Loading ${progress.file}...`,
percent
);
} else if (progress.status === 'done') {
this.updateLoadingStatus('Model loaded successfully!', 100);
}
}
}
);
} catch (error) {
console.log('WebGPU not available, falling back to WASM');
this.generator = await pipeline(
"text-generation",
"onnx-community/Qwen2.5-0.5B-Instruct",
{
dtype: "q8",
progress_callback: (progress) => {
if (progress.status === 'progress') {
const percent = Math.round((progress.loaded / progress.total) * 100);
this.updateLoadingStatus(
`Loading ${progress.file}...`,
percent
);
}
}
}
);
}
}

updateLoadingStatus(message, percent) {
this.elements.loadingStatus.textContent = message;
this.elements.progressFill.style.width = `${percent}%`;
this.elements.progressText.textContent = `${percent}%`;
}

showChatInterface() {
setTimeout(() => {
this.elements.loadingScreen.style.display = 'none';
this.elements.chatInterface.style.display = 'flex';
this.elements.userInput.focus();
}, 500);
}

setupEventListeners() {
// Send button click
this.elements.sendButton.addEventListener('click', () => this.handleSend());

// Enter key to send (Shift+Enter for new line)
this.elements.userInput.addEventListener('keydown', (e) => {
if (e.key === 'Enter' && !e.shiftKey) {
e.preventDefault();
this.handleSend();
}
});

// Auto-resize textarea
this.elements.userInput.addEventListener('input', (e) => {
e.target.style.height = 'auto';
e.target.style.height = Math.min(e.target.scrollHeight, 120) + 'px';
this.elements.sendButton.disabled = !e.target.value.trim() || this.isGenerating;
});

// Suggestion buttons
document.querySelectorAll('.suggestion-btn').forEach(btn => {
btn.addEventListener('click', (e) => {
const prompt = e.target.dataset.prompt;
this.elements.userInput.value = prompt;
this.elements.userInput.dispatchEvent(new Event('input'));
this.handleSend();
});
});
}

async handleSend() {
const message = this.elements.userInput.value.trim();
if (!message || this.isGenerating) return;

// Clear input
this.elements.userInput.value = '';
this.elements.userInput.style.height = 'auto';
this.elements.sendButton.disabled = true;

// Hide welcome message if present
const welcomeMsg = document.querySelector('.welcome-message');
if (welcomeMsg) welcomeMsg.remove();

// Add user message
this.addMessage(message, 'user');

// Add assistant message placeholder with typing indicator
const assistantMessageDiv = this.addMessage('', 'assistant', true);
const messageContent = assistantMessageDiv.querySelector('.message-content');

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

this.isGenerating = true;

try {
let fullResponse = '';

// Create text streamer with callback
const streamer = new TextStreamer(this.generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
fullResponse += text;
// Remove typing indicator if present
const typingIndicator = messageContent.querySelector('.typing-indicator');
if (typingIndicator) typingIndicator.remove();

// Update message content
messageContent.textContent = fullResponse;
this.scrollToBottom();
}
});

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

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

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

// Ensure final text is displayed
messageContent.textContent = finalResponse;

} catch (error) {
console.error('Generation error:', error);
messageContent.textContent = 'Sorry, I encountered an error generating a response. Please try again.';
} finally {
this.isGenerating = false;
this.elements.sendButton.disabled = false;
this.elements.userInput.focus();
this.scrollToBottom();
}
}

addMessage(content, role, showTyping = false) {
const messageDiv = document.createElement('div');
messageDiv.className = `message ${role}-message`;

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

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

if (showTyping) {
const typingIndicator = document.createElement('div');
typingIndicator.className = 'typing-indicator';
typingIndicator.innerHTML = '<span></span><span></span><span></span>';
messageContent.appendChild(typingIndicator);
} else {
messageContent.textContent = content;
}

messageDiv.appendChild(avatar);
messageDiv.appendChild(messageContent);
this.elements.messagesContainer.appendChild(messageDiv);

this.scrollToBottom();
return messageDiv;
}

scrollToBottom() {
this.elements.messagesContainer.scrollTop = this.elements.messagesContainer.scrollHeight;
}

showError(message) {
this.elements.loadingStatus.textContent = message;
this.elements.loadingStatus.style.color = '#FF3B30';
}
}

// Initialize the app
new ChatApp();