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

class Chatbot {
constructor() {
this.generator = null;
this.isGenerating = false;
this.conversationHistory = [];
this.maxHistoryLength = 10;
this.currentModel = 'Xenova/distilgpt2';
this.currentDevice = 'wasm';

this.initializeElements();
this.attachEventListeners();
this.loadModel();
}

initializeElements() {
this.messagesContainer = document.getElementById('messages');
this.messageInput = document.getElementById('message-input');
this.sendButton = document.getElementById('send-button');
this.clearButton = document.getElementById('clear-chat');
this.modelSelect = document.getElementById('model-select');
this.deviceSelect = document.getElementById('device-select');
this.typingIndicator = document.getElementById('typing-indicator');
this.loadingOverlay = document.getElementById('loading-overlay');
this.errorToast = document.getElementById('error-toast');
this.errorMessage = document.getElementById('error-message');
this.charCount = document.getElementById('char-count');
this.dismissErrorBtn = document.getElementById('dismiss-error');
}

attachEventListeners() {
this.sendButton.addEventListener('click', () => this.sendMessage());
this.clearButton.addEventListener('click', () => this.clearChat());
this.modelSelect.addEventListener('change', (e) => this.changeModel(e.target.value));
this.deviceSelect.addEventListener('change', (e) => this.changeDevice(e.target.value));

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

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

this.dismissErrorBtn.addEventListener('click', () => this.hideError());

// Check for WebGPU support
this.checkWebGPUSupport();
}

async checkWebGPUSupport() {
if (!navigator.gpu) {
this.deviceSelect.querySelector('option[value="webgpu"]').disabled = true;
this.deviceSelect.value = 'wasm';
this.currentDevice = 'wasm';
}
}

async loadModel() {
this.showLoading();
try {
const config = {
device: this.currentDevice,
dtype: this.currentDevice === 'webgpu' ? 'fp16' : 'q4'
};

this.generator = await pipeline('text-generation', this.currentModel, config);
this.hideLoading();
this.sendButton.disabled = !this.messageInput.value.trim();
} catch (error) {
console.error('Error loading model:', error);
this.hideLoading();
this.showError('Failed to load model. Please try again or select a different model.');
this.sendButton.disabled = true;
}
}

async changeModel(modelName) {
if (modelName === this.currentModel) return;

this.currentModel = modelName;
await this.loadModel();
}

async changeDevice(device) {
if (device === this.currentDevice) return;

this.currentDevice = device;
await this.loadModel();
}

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

this.addMessage(message, 'user');
this.messageInput.value = '';
this.updateCharCount();
this.autoResizeTextarea();
this.sendButton.disabled = true;
this.isGenerating = true;
this.showTypingIndicator();

try {
const prompt = this.buildPrompt(message);
const response = await this.generateResponse(prompt);
this.addMessage(response, 'bot');
} catch (error) {
console.error('Error generating response:', error);
this.showError('Failed to generate response. Please try again.');
} finally {
this.isGenerating = false;
this.hideTypingIndicator();
this.sendButton.disabled = !this.messageInput.value.trim();
}
}

buildPrompt(userMessage) {
let prompt = "You are a helpful AI assistant. Respond in a friendly and informative way.\n\n";

// Include recent conversation history
if (this.conversationHistory.length > 0) {
const recentHistory = this.conversationHistory.slice(-4);
recentHistory.forEach(msg => {
if (msg.role === 'user') {
prompt += `User: ${msg.content}\n`;
} else {
prompt += `Assistant: ${msg.content}\n`;
}
});
}

prompt += `User: ${userMessage}\nAssistant:`;
return prompt;
}

async generateResponse(prompt) {
const maxNewTokens = 100;
const temperature = 0.7;

const result = await this.generator(prompt, {
max_new_tokens: maxNewTokens,
temperature: temperature,
do_sample: true,
pad_token_id: 50256,
return_full_text: false
});

let response = result[0].generated_text.trim();

// Clean up the response
response = response.replace(/^(Assistant:|AI:|Bot:)/i, '').trim();
response = response.split('\n')[0]; // Take only the first line

if (!response) {
response = "I'm not sure how to respond to that. Could you try rephrasing?";
}

return response;
}

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

const messageContent = document.createElement('div');
messageContent.className = 'message-content';
messageContent.innerHTML = `<p>${this.escapeHtml(content)}</p>`;

const messageTime = document.createElement('div');
messageTime.className = 'message-time';
messageTime.textContent = role === 'user' ? 'You' : 'Bot';

messageDiv.appendChild(messageContent);
messageDiv.appendChild(messageTime);

this.messagesContainer.appendChild(messageDiv);
this.scrollToBottom();

// Update conversation history
this.conversationHistory.push({ role, content });
if (this.conversationHistory.length > this.maxHistoryLength) {
this.conversationHistory.shift();
}
}

clearChat() {
this.messagesContainer.innerHTML = `
<div class="message bot-message">
  <div class="message-content">
    <p>Hello! I'm your AI assistant. How can I help you today?</p>
  </div>
  <div class="message-time">Bot</div>
</div>
`;
this.conversationHistory = [];
this.scrollToBottom();
}

updateCharCount() {
const length = this.messageInput.value.length;
this.charCount.textContent = `${length} / 500`;
}

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

showTypingIndicator() {
this.typingIndicator.classList.remove('hidden');
this.scrollToBottom();
}

hideTypingIndicator() {
this.typingIndicator.classList.add('hidden');
}

showLoading() {
this.loadingOverlay.classList.remove('hidden');
}

hideLoading() {
this.loadingOverlay.classList.add('hidden');
}

showError(message) {
this.errorMessage.textContent = message;
this.errorToast.classList.remove('hidden');
setTimeout(() => this.hideError(), 5000);
}

hideError() {
this.errorToast.classList.add('hidden');
}

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

escapeHtml(text) {
const div = document.createElement('div');
div.textContent = text;
return div.innerHTML;
}
}

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