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

class ChatApp {
constructor() {
this.generator = null;
this.messages = [
{ 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'),
chatMessages: document.getElementById('chatMessages'),
messageInput: document.getElementById('messageInput'),
sendButton: document.getElementById('sendButton'),
progressBar: document.getElementById('progressFill'),
progressText: document.getElementById('progressText')
};

this.init();
}

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

async loadModel() {
try {
this.updateProgress(0, 'Loading model...');

// Create a text generation pipeline with progress tracking
this.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);
this.updateProgress(percent, `Downloading model: ${progress.file}`);
} else if (progress.status === 'done') {
this.updateProgress(100, 'Model loaded successfully!');
}
}
}
);

console.log('Model loaded successfully!');
} catch (error) {
console.error('Model loading error:', error);
throw error;
}
}

updateProgress(percent, text) {
this.elements.progressBar.style.width = `${percent}%`;
this.elements.progressText.textContent = `${percent}%`;

const loadingText = document.querySelector('.loading-text');
if (loadingText) {
loadingText.textContent = text;
}
}

hideLoading() {
this.elements.loadingScreen.classList.add('hidden');
this.elements.messageInput.disabled = false;
this.elements.sendButton.disabled = false;
this.elements.messageInput.focus();
}

setupEventListeners() {
this.elements.sendButton.addEventListener('click', () => this.sendMessage());

this.elements.messageInput.addEventListener('keydown', (e) => {
if (e.key === 'Enter' && !e.shiftKey) {
e.preventDefault();
this.sendMessage();
}
});

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

addWelcomeMessage() {
const welcomeMessage = {
role: 'assistant',
content: "Hello! I'm your AI assistant. How can I help you today? Feel free to ask me anything!" + "};
this.addMessageToChat(welcomeMessage);
}

async sendMessage() {
const userMessage = this.elements.messageInput.value.trim();

if (!userMessage || this.isGenerating) return;

// Add user message to chat
const userMsg = { role: 'user', content: userMessage };
this.messages.push(userMsg);
this.addMessageToChat(userMsg);

// Clear input
this.elements.messageInput.value = '';
this.elements.messageInput.style.height = 'auto';

// Disable input while generating
this.isGenerating = true;
this.elements.messageInput.disabled = true;
this.elements.sendButton.disabled = true;

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

try {
// Generate response with streaming
let fullResponse = '';
let assistantMessageElement = null;

const streamer = new TextStreamer(this.generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
fullResponse += text;

// Remove typing indicator on first token
if (!assistantMessageElement) {
typingIndicator.remove();
assistantMessageElement = this.addMessageToChat({
role: 'assistant',
content: ''
}, true);
}

// Update the message content
const contentElement = assistantMessageElement.querySelector('.message-content');
if (contentElement) {
contentElement.textContent = fullResponse;
this.scrollToBottom();
}
}
});

const output = await this.generator(this.messages, {
max_new_tokens: 512,
do_sample: false,
streamer: streamer,
});

// Add final response to messages history
const assistantResponse = output[0].generated_text.at(-1).content;
this.messages.push({ role: 'assistant', content: assistantResponse });

} catch (error) {
console.error('Generation error:', error);
typingIndicator.remove();
this.addMessageToChat({
role: 'assistant',
content: 'Sorry, I encountered an error. Please try again.'
});
} finally {
this.isGenerating = false;
this.elements.messageInput.disabled = false;
this.elements.sendButton.disabled = false;
this.elements.messageInput.focus();
}
}

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

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

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

messageDiv.appendChild(avatar);
messageDiv.appendChild(content);

this.elements.chatMessages.appendChild(messageDiv);
this.scrollToBottom();

return messageDiv;
}

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

const avatar = document.createElement('div');
avatar.className = 'message-avatar';
avatar.textContent = 'AI';

const content = document.createElement('div');
content.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>';

content.appendChild(typingDiv);
messageDiv.appendChild(avatar);
messageDiv.appendChild(content);

this.elements.chatMessages.appendChild(messageDiv);
this.scrollToBottom();

return messageDiv;
}

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

showError(message) {
const errorDiv = document.createElement('div');
errorDiv.style.cssText = `
position: fixed;
top: 20px;
right: 20px;
background: #ff3b30;
color: white;
padding: 16px 24px;
border-radius: 12px;
box-shadow: 0 4px 12px rgba(0,0,0,0.15);
z-index: 2000;
`;
errorDiv.textContent = message;
document.body.appendChild(errorDiv);

setTimeout(() => errorDiv.remove(), 5000);
}
}

// Initialize the app
new ChatApp();