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

class AppleChatBot {
constructor() {
this.generator = null;
this.isGenerating = false;
this.messages = [
{ role: "system", content: "You are a helpful assistant. Be concise and friendly in your responses." }
];
this.init();
}

async init() {
this.setupEventListeners();
await this.loadModel();
}

setupEventListeners() {
const messageInput = document.getElementById('messageInput');
const sendButton = document.getElementById('sendButton');

// Handle input changes
messageInput.addEventListener('input', () => {
this.updateCharCount();
this.autoResizeTextarea();
sendButton.disabled = !messageInput.value.trim() || this.isGenerating;
});

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

// Handle send button
sendButton.addEventListener('click', () => this.sendMessage());
}

updateCharCount() {
const messageInput = document.getElementById('messageInput');
const charCount = document.getElementById('charCount');
charCount.textContent = `${messageInput.value.length} / 2000`;
}

autoResizeTextarea() {
const textarea = document.getElementById('messageInput');
textarea.style.height = 'auto';
textarea.style.height = Math.min(textarea.scrollHeight, 120) + 'px';
}

async loadModel() {
try {
// Show loading overlay
const loadingOverlay = document.getElementById('loadingOverlay');
const progressFill = document.getElementById('progressFill');
const progressText = document.getElementById('progressText');
const modelStatus = document.getElementById('modelStatus');
const statusDot = modelStatus.querySelector('.status-dot');
const statusText = modelStatus.querySelector('.status-text');

// Update status
statusText.textContent = 'Loading model...';
statusDot.className = 'status-dot loading';

// Create pipeline with progress callback
this.generator = await pipeline(
"text-generation",
"onnx-community/gemma-3-270m-it-ONNX",
{
dtype: "fp32",
progress_callback: (progress) => {
const percentage = Math.round(progress.progress * 100);
progressFill.style.width = `${percentage}%`;
progressText.textContent = `${percentage}%`;

if (progress.status === 'downloading') {
statusText.textContent = `Downloading... ${percentage}%`;
} else if (progress.status === 'ready') {
statusText.textContent = 'Processing model...';
}
}
}
);

// Hide loading overlay
loadingOverlay.style.opacity = '0';
setTimeout(() => {
loadingOverlay.style.display = 'none';
}, 300);

// Update status to ready
statusText.textContent = 'Ready';
statusDot.className = 'status-dot ready';
document.getElementById('sendButton').disabled = false;

} catch (error) {
console.error('Error loading model:', error);
const loadingOverlay = document.getElementById('loadingOverlay');
const loadingContent = loadingOverlay.querySelector('.loading-content');

loadingContent.innerHTML = `
<div class="error-content">
  <i class="fas fa-exclamation-triangle"></i>
  <h2>Failed to Load Model</h2>
  <p>There was an error loading the AI model. Please refresh the page and try again.</p>
  <button onclick="location.reload()" class="retry-button">
                        <i class="fas fa-redo"></i> Retry
                    </button>
</div>
`;
}
}

async sendMessage() {
if (this.isGenerating || !this.generator) return;

const messageInput = document.getElementById('messageInput');
const message = messageInput.value.trim();

if (!message) return;

// Add user message to history
this.messages.push({ role: "user", content: message });
this.addMessageToUI(message, 'user');

// Clear input
messageInput.value = '';
this.updateCharCount();
this.autoResizeTextarea();

// Disable send button
this.isGenerating = true;
document.getElementById('sendButton').disabled = true;

// Show typing indicator
this.showTypingIndicator();

try {
// Create text streamer for real-time output
const streamer = new TextStreamer(this.generator.tokenizer, {
skip_prompt: true,
skip_special_tokens: true,
callback_function: (text) => {
this.updateStreamingMessage(text);
}
});

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

// Get the assistant's response
const assistantMessage = output[0].generated_text.at(-1).content;
this.messages.push({ role: "assistant", content: assistantMessage });

} catch (error) {
console.error('Error generating response:', error);
this.addMessageToUI('Sorry, I encountered an error. Please try again.', 'assistant error');
} finally {
// Hide typing indicator and re-enable send button
this.hideTypingIndicator();
this.isGenerating = false;
document.getElementById('sendButton').disabled = false;
}
}

showTypingIndicator() {
const typingIndicator = document.getElementById('typingIndicator');
typingIndicator.style.display = 'block';

// Create a placeholder message for streaming
this.streamingMessageElement = this.addMessageToUI('', 'assistant', true);
}

hideTypingIndicator() {
const typingIndicator = document.getElementById('typingIndicator');
typingIndicator.style.display = 'none';
}

updateStreamingMessage(text) {
if (this.streamingMessageElement) {
const messageContent = this.streamingMessageElement.querySelector('.message-text');
if (messageContent) {
messageContent.textContent = text;
this.scrollToBottom();
}
}
}

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

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

if (role === 'assistant') {
messageContent.innerHTML = `
<div class="avatar">
  <i class="fas fa-robot"></i>
</div>
<div class="message-content-inner">
  <div class="message-role">AI Assistant</div>
  <div class="message-text">${isStreaming ? '' : this.formatMessage(content)}</div>
</div>
`;
} else if (role === 'user') {
messageContent.innerHTML = `
<div class="avatar user">
  <i class="fas fa-user"></i>
</div>
<div class="message-content-inner">
  <div class="message-role">You</div>
  <div class="message-text">${this.formatMessage(content)}</div>
</div>
`;
} else {
messageContent.innerHTML = `<p>${content}</p>`;
}

messageDiv.appendChild(messageContent);
chatMessages.appendChild(messageDiv);

// Smooth scroll to bottom
this.scrollToBottom();

// Add entrance animation
requestAnimationFrame(() => {
messageDiv.style.opacity = '1';
messageDiv.style.transform = 'translateY(0)';
});

return messageDiv;
}

formatMessage(content) {
// Basic markdown-like formatting
return content
.replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
.replace(/\*(.*?)\*/g, '<em>$1</em>')
.replace(/`(.*?)`/g, '<code>$1</code>')
.replace(/\n/g, '<br>');
    }

    scrollToBottom() {
        const chatMessages = document.getElementById('chatMessages');
        chatMessages.scrollTop = chatMessages.scrollHeight;
    }
}

// Initialize the chatbot when the page loads
document.addEventListener('DOMContentLoaded', () => {
    new AppleChatBot();
});