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

// Configure environment
env.allowLocalModels = false;

// State
let generator = null;
let currentModel = null;
let isVisionModel = false;
let currentImage = null;
let isGenerating = false;
let conversationHistory = [];

// DOM Elements
const modelSelect = document.getElementById('model-select');
const loadModelBtn = document.getElementById('load-model-btn');
const loadingContainer = document.getElementById('loading-container');
const loadingStatus = document.getElementById('loading-status');
const progressBar = document.getElementById('progress-bar');
const progressText = document.getElementById('progress-text');
const chatContainer = document.getElementById('chat-container');
const chatMessages = document.getElementById('chat-messages');
const imageInput = document.getElementById('image-input');
const attachImageBtn = document.getElementById('attach-image-btn');
const imageUrlInput = document.getElementById('image-url-input');
const loadUrlBtn = document.getElementById('load-url-btn');
const imagePreviewContainer = document.getElementById('image-preview-container');
const imagePreview = document.getElementById('image-preview');
const removeImageBtn = document.getElementById('remove-image-btn');
const userInput = document.getElementById('user-input');
const sendBtn = document.getElementById('send-btn');
const errorContainer = document.getElementById('error-container');
const errorMessage = document.getElementById('error-message');
const dismissErrorBtn = document.getElementById('dismiss-error-btn');

// Settings
const maxTokensSlider = document.getElementById('max-tokens');
const maxTokensValue = document.getElementById('max-tokens-value');
const temperatureSlider = document.getElementById('temperature');
const temperatureValue = document.getElementById('temperature-value');
const topPSlider = document.getElementById('top-p');
const topPValue = document.getElementById('top-p-value');

// Vision model identifiers
const VISION_MODELS = ['SmolVLM', 'Fara', 'llava', 'vision'];

function isVisionModelSelected(modelId) {
return VISION_MODELS.some(vm => modelId.toLowerCase().includes(vm.toLowerCase()));
}

// Progress callback for model loading
function progressCallback(progress) {
if (progress.status === 'initiate') {
loadingStatus.textContent = `Loading ${progress.file || 'model'}...`;
} else if (progress.status === 'download') {
loadingStatus.textContent = `Downloading ${progress.file || 'model'}...`;
} else if (progress.status === 'progress') {
const percent = Math.round(progress.progress || 0);
progressBar.style.width = `${percent}%`;
progressText.textContent = `${percent}%`;
loadingStatus.textContent = `Downloading ${progress.file || 'model'}...`;
} else if (progress.status === 'done') {
loadingStatus.textContent = `Loaded ${progress.file || 'model'}`;
} else if (progress.status === 'ready') {
loadingStatus.textContent = 'Model ready!';
progressBar.style.width = '100%';
progressText.textContent = '100%';
}
}

// Load model
async function loadModel() {
const modelId = modelSelect.value;

if (currentModel === modelId && generator) {
showError('Model already loaded!');
return;
}

try {
loadModelBtn.disabled = true;
loadingContainer.classList.remove('hidden');
chatContainer.classList.add('hidden');
progressBar.style.width = '0%';
progressText.textContent = '0%';
loadingStatus.textContent = 'Initializing...';

isVisionModel = isVisionModelSelected(modelId);

// Clean up previous model
if (generator) {
generator = null;
}

// Determine device - try WebGPU first
let device = 'wasm';
if (navigator.gpu) {
try {
const adapter = await navigator.gpu.requestAdapter();
if (adapter) {
device = 'webgpu';
loadingStatus.textContent = 'Using WebGPU acceleration...';
}
} catch (e) {
console.log('WebGPU not available, falling back to WASM');
}
}

loadingStatus.textContent = `Loading model on ${device.toUpperCase()}...`;

// Create pipeline based on model type
if (isVisionModel) {
generator = await pipeline('image-text-to-text', modelId, {
device: device,
dtype: 'q4f16',
progress_callback: progressCallback,
});
} else {
generator = await pipeline('text-generation', modelId, {
device: device,
dtype: 'q4f16',
progress_callback: progressCallback,
});
}

currentModel = modelId;
conversationHistory = [];

// Update UI
loadingContainer.classList.add('hidden');
chatContainer.classList.remove('hidden');
sendBtn.disabled = false;

// Update attach button visibility
attachImageBtn.style.display = isVisionModel ? 'block' : 'none';
imageUrlInput.style.display = isVisionModel ? 'block' : 'none';
loadUrlBtn.style.display = isVisionModel ? 'block' : 'none';

// Clear chat and show ready message
chatMessages.innerHTML = `
<div class="welcome-message">
  <p>✅ <strong>${modelId.split('/').pop()}</strong> loaded successfully!</p>
  <p class="hint">${isVisionModel ? 'This is a vision model. You can attach images to your messages.' : 'This is a
    text-only model for conversation.'}</p>
</div>
`;

} catch (error) {
console.error('Error loading model:', error);
showError(`Failed to load model: ${error.message}`);
loadingContainer.classList.add('hidden');
} finally {
loadModelBtn.disabled = false;
}
}

// Handle image upload
function handleImageUpload(file) {
if (!file) return;

const reader = new FileReader();
reader.onload = (e) => {
currentImage = e.target.result;
imagePreview.src = currentImage;
imagePreviewContainer.classList.remove('hidden');
};
reader.readAsDataURL(file);
}

// Load image from URL
async function loadImageFromUrl() {
const url = imageUrlInput.value.trim();
if (!url) return;

try {
currentImage = url;
imagePreview.src = url;
imagePreviewContainer.classList.remove('hidden');
imageUrlInput.value = '';
} catch (error) {
showError('Failed to load image from URL');
}
}

// Remove attached image
function removeImage() {
currentImage = null;
imagePreview.src = '';
imagePreviewContainer.classList.add('hidden');
imageInput.value = '';
}

// Add message to chat
function addMessage(role, content, imageUrl = null) {
const messageDiv = document.createElement('div');
messageDiv.className = `message ${role}`;

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

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

if (imageUrl) {
const img = document.createElement('img');
img.src = imageUrl;
img.className = 'message-image';
contentDiv.appendChild(img);
}

const textSpan = document.createElement('span');
textSpan.textContent = content;
contentDiv.appendChild(textSpan);

messageDiv.appendChild(label);
messageDiv.appendChild(contentDiv);
chatMessages.appendChild(messageDiv);

chatMessages.scrollTop = chatMessages.scrollHeight;

return contentDiv;
}

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

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

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

const typingDiv = document.createElement('div');
typingDiv.className = 'typing-indicator';
typingDiv.innerHTML = '<span></span><span></span><span></span>';

contentDiv.appendChild(typingDiv);
messageDiv.appendChild(label);
messageDiv.appendChild(contentDiv);
chatMessages.appendChild(messageDiv);

chatMessages.scrollTop = chatMessages.scrollHeight;
}

// Remove typing indicator
function removeTypingIndicator() {
const indicator = document.getElementById('typing-indicator');
if (indicator) {
indicator.remove();
}
}

// Send message
async function sendMessage() {
const text = userInput.value.trim();
if (!text || !generator || isGenerating) return;

isGenerating = true;
sendBtn.disabled = true;
userInput.disabled = true;

try {
// Add user message
addMessage('user', text, currentImage);

// Clear input
userInput.value = '';
const imageForMessage = currentImage;
removeImage();

// Show typing indicator
addTypingIndicator();

// Get generation settings
const maxTokens = parseInt(maxTokensSlider.value);
const temperature = parseFloat(temperatureSlider.value);
const topP = parseFloat(topPSlider.value);

let response;

if (isVisionModel && imageForMessage) {
// Vision model with image
const messages = [
{
role: 'user',
content: [
{ type: 'image', image: imageForMessage },
{ type: 'text', text: text }
]
}
];

const output = await generator(messages, {
max_new_tokens: maxTokens,
temperature: temperature,
top_p: topP,
do_sample: temperature > 0,
});

response = output[0].generated_text.at(-1).content;

} else if (isVisionModel) {
// Vision model without image - text only
const messages = [
{
role: 'user',
content: [
{ type: 'text', text: text }
]
}
];

const output = await generator(messages, {
max_new_tokens: maxTokens,
temperature: temperature,
top_p: topP,
do_sample: temperature > 0,
});

response = output[0].generated_text.at(-1).content;

} else {
// Text-only model
conversationHistory.push({ role: 'user', content: text });

const output = await generator(conversationHistory, {
max_new_tokens: maxTokens,
temperature: temperature,
top_p: topP,
do_sample: temperature > 0,
});

const generatedMessages = output[0].generated_text;
const assistantMessage = generatedMessages[generatedMessages.length - 1];
response = assistantMessage.content;

conversationHistory.push({ role: 'assistant', content: response });
}

// Remove typing indicator and add response
removeTypingIndicator();
addMessage('assistant', response);

} catch (error) {
console.error('Error generating response:', error);
removeTypingIndicator();
showError(`Generation error: ${error.message}`);
} finally {
isGenerating = false;
sendBtn.disabled = false;
userInput.disabled = false;
userInput.focus();
}
}

// Show error
function showError(message) {
errorMessage.textContent = message;
errorContainer.classList.remove('hidden');
}

// Hide error
function hideError() {
errorContainer.classList.add('hidden');
}

// Event Listeners
loadModelBtn.addEventListener('click', loadModel);

attachImageBtn.addEventListener('click', () => {
if (isVisionModel) {
imageInput.click();
}
});

imageInput.addEventListener('change', (e) => {
if (e.target.files.length > 0) {
handleImageUpload(e.target.files[0]);
}
});

loadUrlBtn.addEventListener('click', loadImageFromUrl);

imageUrlInput.addEventListener('keypress', (e) => {
if (e.key === 'Enter') {
loadImageFromUrl();
}
});

removeImageBtn.addEventListener('click', removeImage);

sendBtn.addEventListener('click', sendMessage);

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

dismissErrorBtn.addEventListener('click', hideError);

// Settings sliders
maxTokensSlider.addEventListener('input', () => {
maxTokensValue.textContent = maxTokensSlider.value;
});

temperatureSlider.addEventListener('input', () => {
temperatureValue.textContent = temperatureSlider.value;
});

topPSlider.addEventListener('input', () => {
topPValue.textContent = topPSlider.value;
});

// Drag and drop for images
chatContainer.addEventListener('dragover', (e) => {
if (isVisionModel) {
e.preventDefault();
e.dataTransfer.dropEffect = 'copy';
}
});

chatContainer.addEventListener('drop', (e) => {
if (isVisionModel) {
e.preventDefault();
const files = e.dataTransfer.files;
if (files.length > 0 && files[0].type.startsWith('image/')) {
handleImageUpload(files[0]);
}
}
});

// Initialize
document.addEventListener('DOMContentLoaded', () => {
// Hide image controls initially
attachImageBtn.style.display = 'none';
imageUrlInput.style.display = 'none';
loadUrlBtn.style.display = 'none';
});