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

class ChatApp {
constructor() {
this.generator = null;
this.messages = [];
this.isGenerating = false;
this.currentStreamingMessage = null;

this.initElements();
this.initEventListeners();
this.initWorker();
}

initElements() {
this.messagesContainer = document.getElementById('messages');
this.messageInput = document.getElementById('messageInput');
this.sendButton = document.getElementById('sendButton');
this.statusText = document.getElementById('statusText');
this.loadingOverlay = document.getElementById('loadingOverlay');
this.loadingProgress = document.getElementById('loadingProgress');
this.progressFill = document.getElementById('progressFill');
this.welcomeMessage = document.querySelector('.welcome-message');
}

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

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

// Example prompt chips
document.querySelectorAll('.example-chip').forEach(chip => {
chip.addEventListener('click', () => {
const prompt = chip.dataset.prompt;
this.messageInput.value = prompt;
this.messageInput.focus();
this.messageInput.style.height = 'auto';
this.messageInput.style.height = Math.min(this.messageInput.scrollHeight, 120) + 'px';
});
});
}

initWorker() {
// Create a worker to load the model in the background
const workerCode = `
self.addEventListener('message', async function(e) {
if (e.data.type === 'load') {
try {
const { pipeline } = await import('https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.1');

const generator = await pipeline(
"text-generation",
"onnx-community/gemma-3-270m-it-ONNX",
{
dtype: "fp32",
progress_callback: (info) => {
self.postMessage({
type: 'progress',
progress: info.progress || 0,
status: info.status || 'Loading...'
});
}
}
);

self.postMessage({ type: 'ready' });
} catch (error) {
self.postMessage({
type: 'error',
error: error.message
});
}
}
});
`;

const blob = new Blob([workerCode], { type: 'application/javascript' });
this.worker = new Worker(URL.createObjectURL(blob));

this.worker.addEventListener('message', (e) => {
switch (e.data.type) {
case 'progress':
this.updateProgress(e.data.progress, e.data.status);
break;
case 'ready':
this.onModelReady();
break;
case 'error':
this.onError(e.data.error);
break;
}
});

this.worker.postMessage({ type: 'load' });
}

updateProgress(progress, status) {
this.progressFill.style.width = `${Math.round(progress * 100)}%`;
this.loadingProgress.textContent = status;
}

async onModelReady() {
try {
// Load the pipeline in the main thread
this.generator = await pipeline(
"text-generation",
"onnx-community/gemma-3-270m-it-ONNX",
{
dtype: "fp32",
progress_callback: (info) => {
if (info.status) {
this.statusText.textContent = info.status;
}
}
}
);

this.loadingOverlay.style.display = 'none';
this.messageInput.disabled = false;
this.sendButton.disabled = false;
this.statusText.textContent = 'Ready to chat';
this.messageInput.focus();
} catch (error) {
this.onError(error.message);
}
}

onError(error) {
this.loadingProgress.textContent = 'Error loading model';
this.statusText.textContent = 'Error: ' + error;
console.error('Model loading error:', error);
}

async sendMessage() {
const text = this.messageInput.value.trim();
if (!text || this.isGenerating) return;

// Hide welcome message on first interaction
if (this.welcomeMessage) {
this.welcomeMessage.style.display = 'none';
}

// Add user message
this.addMessage('user', text);
this.messageInput.value = '';
this.messageInput.style.height = 'auto';

// Add assistant message placeholder
const assistantMessageId = this.addMessage('assistant', '', true);

this.isGenerating = true;
this.sendButton.disabled = true;
this.statusText.textContent = 'Generating response...';

try {
// Prepare messages for the model
const modelMessages = [
{ role: "system", content: "You are a helpful, friendly, and knowledgeable assistant. Provide clear, concise, and
accurate responses." },
...this.messages.map(msg => ({
role: msg.role === 'user' ? 'user' : 'assistant',
content: msg.content
})),
{ role: "user", content: text }
];

// Create a custom streamer to capture the text
let streamedText = '';
const streamer = new TextStreamer(this.generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
streamedText += text;
this.updateMessage(assistantMessageId, streamedText);
}
});

// Generate response with streaming
await this.generator(modelMessages, {
max_new_tokens: 512,
do_sample: false,
temperature: 0.7,
streamer: streamer
});

// Save the complete message
this.messages.push({ role: 'assistant', content: streamedText });

} catch (error) {
console.error('Generation error:', error);
this.updateMessage(assistantMessageId, 'Sorry, I encountered an error. Please try again.');
} finally {
this.isGenerating = false;
this.sendButton.disabled = false;
this.statusText.textContent = 'Ready to chat';
this.messageInput.focus();
}
}

addMessage(role, content, isTyping = false) {
const messageId = 'msg-' + Date.now();
const messageDiv = document.createElement('div');
messageDiv.className = `message ${role}`;
messageDiv.id = messageId;

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

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

const textDiv = document.createElement('div');
textDiv.className = 'message-text';

if (isTyping) {
textDiv.innerHTML = '<div class="typing-indicator">
  <div class="typing-dot"></div>
  <div class="typing-dot"></div>
  <div class="typing-dot"></div>
</div>';
} else {
textDiv.textContent = content;
}

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

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

if (!isTyping && role === 'user') {
this.messages.push({ role, content });
}

return messageId;
}

updateMessage(messageId, content) {
const messageElement = document.getElementById(messageId);
if (messageElement) {
const textElement = messageElement.querySelector('.message-text');
textElement.textContent = content;
this.messagesContainer.scrollTop = this.messagesContainer.scrollHeight;
}
}
}

// Initialize the app when DOM is loaded
document.addEventListener('DOMContentLoaded', () => {
new ChatApp();
});