// ============================================
// LM Studio Web Chat — avec génération d'image
// ============================================

// Configuration
const MCP_SERVER_URL = 'http://localhost:8085';  // Backend FastAPI pour ComfyUI

// Presets serveurs
const SERVER_PRESETS = {
  ollama: 'http://localhost:11434',
  lmstudio: 'http://localhost:1234',
  openrouter: 'https://openrouter.ai/api/v1'
};

// Détection de commande /genimg
const GENIMG_PATTERN = /^\/genimg\s+(.+)$/i;

// Variables d'état
let isConnected = false;
let currentModel = '';
let pendingImage = null;
let systemPrompt = '';

// Chat management
let chats = [];
let currentChat = null;

// Définition des éléments DOM
const chatContainer = document.getElementById('chat-container');
const userInput = document.getElementById('user-input');
const serverTypeSelect = document.getElementById('server-type-select');
const customUrlInput = document.getElementById('custom-url');
const apiKeyInput = document.getElementById('api-key');
const serverUrlInput = document.getElementById('server-url-input');
const connectButton = document.getElementById('connect-button');
const connectionStatus = document.getElementById('connection-status');
const sendButton = document.getElementById('send-button');
const newChatButton = document.getElementById('new-chat-button');
const toggleSidebarButton = document.getElementById('toggle-sidebar');
const chatSidebar = document.getElementById('chat-sidebar');
const chatList = document.getElementById('chat-list');
const contextMenu = document.getElementById('context-menu');
const modelSelect = document.getElementById('model-select');
const uploadButton = document.getElementById('upload-button');
const imageUpload = document.getElementById('image-upload');
const imagePreview = document.getElementById('image-preview');
const systemPromptInput = document.getElementById('system-prompt-input');
const resetSystemPromptButton = document.getElementById('reset-system-prompt');

// ============================================
// Initialisation System Prompt
// ============================================
function initSystemPrompt() {
  const savedPrompt = localStorage.getItem('lmstudio-system-prompt');
  const defaultPrompt = "You are an intelligent assistant. You always provide well-reasoned answers that are both correct and helpful.";
  systemPrompt = savedPrompt || defaultPrompt;
  systemPromptInput.value = systemPrompt;
}

systemPromptInput.addEventListener('input', () => {
  systemPrompt = systemPromptInput.value;
  localStorage.setItem('lmstudio-system-prompt', systemPrompt);
});

resetSystemPromptButton.addEventListener('click', () => {
  const defaultPrompt = "You are an intelligent assistant. You always provide well-reasoned answers that are both correct and helpful.";
  systemPrompt = defaultPrompt;
  systemPromptInput.value = systemPrompt;
  localStorage.setItem('lmstudio-system-prompt', systemPrompt);
});

// ============================================
// Sélection de Serveur et API Key
// ============================================
function setupServerSelection() {
  const savedServerType = localStorage.getItem('server-type') || 'ollama';
  serverTypeSelect.value = savedServerType;
  updateServerUI();

  serverTypeSelect.addEventListener('change', () => {
    localStorage.setItem('server-type', serverTypeSelect.value);
    updateServerUI();
    if (serverTypeSelect.value && SERVER_PRESETS[serverTypeSelect.value]) {
      serverUrlInput.value = SERVER_PRESETS[serverTypeSelect.value];
    }
  });

  customUrlInput.addEventListener('input', () => {
    if (serverTypeSelect.value === 'custom') {
      serverUrlInput.value = customUrlInput.value.trim();
    }
  });

  apiKeyInput.addEventListener('keypress', (e) => {
    if (e.key === 'Enter') {
      userInput.focus();
    }
  });

  apiKeyInput.addEventListener('input', () => {
    localStorage.setItem('api-key', apiKeyInput.value);
  });

  const savedApiKey = localStorage.getItem('api-key');
  if (savedApiKey) {
    apiKeyInput.value = savedApiKey;
  }
}

function updateServerUI() {
  const type = serverTypeSelect.value;
  if (!type) {
    customUrlInput.style.display = 'none';
    apiKeyInput.style.display = 'none';
    serverUrlInput.value = '';
  } else if (type === 'custom') {
    customUrlInput.style.display = 'block';
    apiKeyInput.style.display = 'block';
    serverUrlInput.value = customUrlInput.value.trim();
  } else {
    const url = SERVER_PRESETS[type];
    serverUrlInput.value = url;
    customUrlInput.style.display = 'none';
    apiKeyInput.style.display = (type === 'openrouter') ? 'block' : 'none';
  }
}

// ============================================
// Connexion au Serveur LLM
// ============================================
async function connectToServer() {
  const serverUrl = serverUrlInput.value.trim();
  if (!serverUrl) {
    updateConnectionStatus('Please enter a valid server address', false);
    return;
  }
  const apiKey = apiKeyInput.value.trim();
  try {
    updateConnectionStatus('Connecting...', false);
    const headers = { 'Content-Type': 'application/json' };
    if (apiKey) {
      headers['Authorization'] = `Bearer ${apiKey}`;
    }
    const response = await fetch(`${serverUrl}/v1/models`, {
      method: 'GET',
      headers: headers
    });
    if (!response.ok) throw new Error('Server response was not ok');
    const data = await response.json();
    if (data && data.data && data.data.length > 0) {
      modelSelect.innerHTML = "";
      data.data.forEach(model => {
        const option = document.createElement('option');
        option.value = model.id;
        option.textContent = model.id;
        modelSelect.appendChild(option);
      });
      modelSelect.disabled = false;
      // S'assurer qu'une option est sélectionnée
      if (modelSelect.options.length > 0) {
        modelSelect.selectedIndex = 0;
      }
      currentModel = modelSelect.value;
      isConnected = true;
      updateConnectionStatus('Connected', true);
      userInput.disabled = false;
      sendButton.disabled = false;
      if (!currentChat) createNewChat();
      const serverLabel = serverTypeSelect.value ? ` ${serverTypeSelect.options[serverTypeSelect.selectedIndex].text}` : '';
      addMessage(`Connected to${serverLabel}. You can start chatting now!`, false, null, false);
    } else {
      throw new Error('No models available');
    }
  } catch (error) {
    console.error('Error:', error);
    updateConnectionStatus('Failed to connect', false);
    addMessage('Error: Unable to connect to the server. Please check the address and try again.', false);
  }
}

function updateConnectionStatus(message, connected) {
  connectionStatus.textContent = message;
  connectionStatus.style.color = connected ? 'var(--accent-color)' : '#f44336';
  connectButton.textContent = connected ? 'Disconnect' : 'Connect';
  customUrlInput.disabled = connected;
  apiKeyInput.disabled = connected;
  serverTypeSelect.disabled = connected;
  userInput.disabled = !connected;
  sendButton.disabled = !connected;
}

// ============================================
// Gestion des Chats
// ============================================
function createNewChat() {
  const chatId = Date.now();
  const newChat = { id: chatId, name: `Conversation ${chats.length + 1}`, messages: [] };
  chats.push(newChat);
  currentChat = newChat;
  updateChatList();
  chatContainer.innerHTML = '';
}

function updateChatList() {
  chatList.innerHTML = '';
  chats.forEach(chat => {
    const li = document.createElement('li');
    li.textContent = chat.name;
    li.dataset.chatId = chat.id;
    if (currentChat && chat.id === currentChat.id) li.classList.add('active');
    li.addEventListener('click', () => {
      if (currentChat && chat.id === currentChat.id) return;
      currentChat = chat;
      loadChat(chat);
      updateChatList();
    });
    li.addEventListener('contextmenu', (e) => {
      e.preventDefault();
      showContextMenu(e.pageX, e.pageY, chat.id);
    });
    chatList.appendChild(li);
  });
}

function loadChat(chat) {
  chatContainer.innerHTML = '';
  chat.messages.forEach(message => {
    addMessage(message.content, message.isUser, message.metrics, false);
  });
}

function showContextMenu(x, y, chatId) {
  contextMenu.style.left = x + "px";
  contextMenu.style.top = y + "px";
  contextMenu.style.display = "block";
  contextMenu.onclick = () => {
    deleteChat(chatId);
    hideContextMenu();
  };
}
function hideContextMenu() { contextMenu.style.display = "none"; }
document.addEventListener('click', () => { if (contextMenu.style.display === "block") hideContextMenu(); });

function deleteChat(chatId) {
  chats = chats.filter(c => c.id != chatId);
  if (currentChat && currentChat.id == chatId) {
    currentChat = chats.length > 0 ? chats[0] : null;
    if (!currentChat) createNewChat();
    else loadChat(currentChat);
  }
  updateChatList();
}

// ============================================
// Ejection de Modèle
// ============================================
async function ejectCurrentModel(oldModel) {
  const serverUrl = serverUrlInput.value.trim();
  const apiKey = apiKeyInput.value.trim();
  try {
    const headers = { 'Content-Type': 'application/json' };
    if (apiKey) {
      headers['Authorization'] = `Bearer ${apiKey}`;
    }
    await fetch(`${serverUrl}/v1/model/eject`, {
      method: 'POST',
      headers: headers,
      body: JSON.stringify({ model: oldModel })
    });
    console.log(`Model ${oldModel} ejected.`);
  } catch (error) {
    console.error("Error ejecting model:", error);
  }
}

modelSelect.addEventListener('change', async (e) => {
  const newModel = e.target.value;
  if (currentModel && currentModel !== newModel) {
    modelSelect.disabled = true;
    await ejectCurrentModel(currentModel);
  }
  currentModel = newModel;
  modelSelect.disabled = false;
});

// ============================================
// Construction de l'Historique
// ============================================
function buildConversationHistory() {
  let effectiveSystemPrompt = systemPrompt;
  if (currentChat.messages.some(msg => msg.isImage) && !systemPrompt.trim()) {
    effectiveSystemPrompt = "You are an AI assistant that analyzes images.";
  } else if (!systemPrompt.trim()) {
    effectiveSystemPrompt = "You are an intelligent assistant. You always provide well-reasoned answers that are both correct and helpful.";
  }
  const history = [{ role: 'system', content: effectiveSystemPrompt }];
  currentChat.messages.forEach(msg => {
    if (msg.isImage) {
      // Pour les modèles ne supportant pas la vision, on envoie seulement le texte
      // Le modèle ne verra pas l'image mais l'historique reste valide
      const textContent = msg.text || "Image generated.";
      history.push({
        role: msg.isUser ? 'user' : 'assistant',
        content: textContent
      });
    } else {
      history.push({ role: msg.isUser ? 'user' : 'assistant', content: msg.content });
    }
  });
  return history;
}

// ============================================
// Ajout de Messages (texte & image)
// ============================================
function addMessage(content, isUser, metrics = null, store = true) {
  const messageDiv = document.createElement('div');
  messageDiv.classList.add('message', isUser ? 'user-message' : 'assistant-message');

  const headerDiv = document.createElement('div');
  headerDiv.classList.add('message-header');
  headerDiv.textContent = isUser ? 'You' : 'Assistant';
  messageDiv.appendChild(headerDiv);

  if (!isUser && currentModel) {
    const modelDiv = document.createElement('div');
    modelDiv.classList.add('message-model');
    modelDiv.textContent = currentModel;
    messageDiv.appendChild(modelDiv);
  }

  const contentDiv = document.createElement('div');
  contentDiv.classList.add('message-content');
  contentDiv.innerHTML = marked.parse(content);
  messageDiv.appendChild(contentDiv);

  if (metrics) {
    const metricsDiv = document.createElement('div');
    metricsDiv.classList.add('message-metrics');
    metricsDiv.textContent = metrics;
    messageDiv.appendChild(metricsDiv);
  }

  chatContainer.appendChild(messageDiv);
  chatContainer.scrollTop = chatContainer.scrollHeight;

  messageDiv.querySelectorAll('pre code').forEach(block => {
    hljs.highlightElement(block);
  });

  if (typeof MathJax !== 'undefined') {
    MathJax.typesetPromise([messageDiv]).catch(err => console.error('MathJax typeset failed:', err));
  }

  attachCopyButtons(messageDiv);

  if (store && currentChat) {
    currentChat.messages.push({ content, isUser, metrics, isImage: false });
  }
}

// ============================================
// Boutons de Copie pour Code Blocks
// ============================================
function attachCopyButtons(container) {
  container.querySelectorAll('pre').forEach(pre => {
    if (pre.querySelector('.copy-btn')) return; // déjà un bouton

    const btn = document.createElement('button');
    btn.className = 'copy-btn';
    btn.textContent = 'Copy';
    btn.onclick = () => {
      const code = pre.querySelector('code');
      const text = code ? code.textContent : pre.textContent;
      navigator.clipboard.writeText(text).then(() => {
        btn.textContent = 'Copied!';
        setTimeout(() => btn.textContent = 'Copy', 2000);
      });
    };
    pre.style.position = 'relative';
    pre.appendChild(btn);
  });
}

// Lightbox functions
function openLightbox(imageUrl) {
  const lightbox = document.getElementById('lightbox');
  const lightboxImg = document.getElementById('lightbox-image');
  if (lightbox && lightboxImg) {
    lightboxImg.src = imageUrl;
    lightbox.classList.add('active');
    document.body.style.overflow = 'hidden';
  }
}

function closeLightbox() {
  const lightbox = document.getElementById('lightbox');
  if (lightbox) {
    lightbox.classList.remove('active');
    document.body.style.overflow = '';
  }
}

document.addEventListener('DOMContentLoaded', () => {
  const lightbox = document.getElementById('lightbox');
  const closeBtn = document.querySelector('.lightbox-close');
  if (lightbox) {
    lightbox.addEventListener('click', (e) => {
      if (e.target === lightbox) closeLightbox();
    });
    if (closeBtn) {
      closeBtn.addEventListener('click', closeLightbox);
    }
  }
});

function addImageMessage(dataURL, promptText) {
  addMessage(`<img src="${dataURL}" class="generated-image" alt="Generated image" />`, false);
  const lastMsgDiv = chatContainer.lastElementChild;
  const img = lastMsgDiv.querySelector('.generated-image');
  if (img) {
    img.addEventListener('click', () => openLightbox(dataURL));
  }
  const lastMsg = currentChat.messages[currentChat.messages.length - 1];
  lastMsg.isImage = true;
  lastMsg.imageData = dataURL;
  lastMsg.text = promptText;
}

// ============================================
// Génération d'Image via /genimg
// ============================================
async function handleImageGeneration(description) {
  addMessage(`/genimg ${description}`, true);

  const generatingDiv = document.createElement('div');
  generatingDiv.classList.add('message', 'assistant-message');
  generatingDiv.innerHTML = `
    <div class="message-header">Assistant</div>
    <div class="message-content">
      <p><i class="fas fa-spinner fa-spin"></i> Génération d'image en cours... (peut prendre jusqu'à 60 secondes)</p>
    </div>
  `;
  chatContainer.appendChild(generatingDiv);
  chatContainer.scrollTop = chatContainer.scrollHeight;

  userInput.disabled = true;
  sendButton.disabled = true;

  try {
    const response = await fetch(`${MCP_SERVER_URL}/generate-image`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        positive_prompt: description,
        negative_prompt: 'text, watermark, blurry, low quality'
      })
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({}));
      throw new Error(errorData.detail || `HTTP ${response.status}: ${response.statusText}`);
    }

    const data = await response.json();
    generatingDiv.remove();
    addImageMessage(data.image_url, description);

  } catch (error) {
    console.error('Image generation error:', error);
    generatingDiv.remove();
    addMessage(`Erreur lors de la génération d'image: ${error.message}`, false);
  } finally {
    userInput.disabled = false;
    sendButton.disabled = false;
    userInput.focus();
  }
}

// ============================================
// Envoi de Messages (Chat & Commande /genimg)
// ============================================
async function sendMessage() {
  let message = userInput.value.trim();
  if (!message && !pendingImage) return;

  // Détection /genimg
  const genimgMatch = message.match(GENIMG_PATTERN);
  if (genimgMatch) {
    const description = genimgMatch[1].trim();
    userInput.value = '';
    await handleImageGeneration(description);
    return;
  }

  if (pendingImage) {
    let promptText = message || "What's in this image?";
    addImageMessage(pendingImage, promptText);
    pendingImage = null;
    imagePreview.style.display = "none";
    userInput.value = "";
  } else {
    addMessage(message, true);
  }

  const conversationHistory = buildConversationHistory();

  // Message assistant temporaire pour streaming
  const assistantMessageElement = document.createElement('div');
  assistantMessageElement.classList.add('message', 'assistant-message');

  const headerDiv = document.createElement('div');
  headerDiv.classList.add('message-header');
  headerDiv.textContent = 'Assistant';
  assistantMessageElement.appendChild(headerDiv);

  if (currentModel) {
    const modelDiv = document.createElement('div');
    modelDiv.classList.add('message-model');
    modelDiv.textContent = currentModel;
    assistantMessageElement.appendChild(modelDiv);
  }

  const assistantContentDiv = document.createElement('div');
  assistantContentDiv.classList.add('message-content');
  assistantContentDiv.innerHTML = '<p><i class="fas fa-spinner fa-spin"></i> Thinking...</p>';
  assistantMessageElement.appendChild(assistantContentDiv);

  chatContainer.appendChild(assistantMessageElement);
  chatContainer.scrollTop = chatContainer.scrollHeight;

  userInput.value = '';
  userInput.disabled = true;
  sendButton.disabled = true;

  const serverUrl = serverUrlInput.value.trim();
  const startTime = performance.now();
  let accumulatedText = '';

  try {
    const headers = { 'Content-Type': 'application/json' };
    const apiKey = apiKeyInput.value.trim();
    if (apiKey) {
      headers['Authorization'] = `Bearer ${apiKey}`;
    }
    const response = await fetch(`${serverUrl}/v1/chat/completions`, {
      method: 'POST',
      headers: headers,
      body: JSON.stringify({
        model: currentModel,
        messages: conversationHistory,
        temperature: 0.7,
        max_tokens: 4096,
        stream: true
      })
    });
    if (!response.ok) throw new Error('Server response was not ok');

    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let done = false;
    while (!done) {
      const { value, done: doneReading } = await reader.read();
      done = doneReading;
      if (value) {
        const chunk = decoder.decode(value, { stream: true });
        const lines = chunk.split('\n').filter(line => line.trim() !== '');
        for (const line of lines) {
          if (line.startsWith("data:")) {
            const dataStr = line.slice(5).trim();
            if (dataStr === "[DONE]") { done = true; break; }
            try {
              const parsed = JSON.parse(dataStr);
              const delta = parsed.choices[0].delta;
              if (delta && delta.content) {
                accumulatedText += delta.content;
                assistantContentDiv.innerHTML = marked.parse(accumulatedText);
                assistantMessageElement.querySelectorAll('pre code').forEach(block => {
                  hljs.highlightElement(block);
                });
                attachCopyButtons(assistantMessageElement);
                if (typeof MathJax !== 'undefined') {
                  MathJax.typesetPromise([assistantMessageElement]).catch(err => console.error(err));
                }
              }
            } catch (err) {
              console.error("Error parsing stream chunk", err);
            }
          } else if (line.startsWith("event:")) {
            const eventType = line.slice(6).trim();
            if (eventType === "error") {
              console.error("Received error event from server:", line);
              addMessage("Error: Received error event from server", false);
              done = true;
              break;
            }
          }
        }
      }
    }
    const endTime = performance.now();
    const timeElapsed = ((endTime - startTime) / 1000).toFixed(2);
    if (currentChat) {
      currentChat.messages.push({ content: accumulatedText, isUser: false, isImage: false });
      if (currentChat.name.startsWith('Conversation')) {
        const snippet = accumulatedText.split(' ').slice(0, 7).join(' ');
        currentChat.name = snippet ? `Conversation: ${snippet}...` : currentChat.name;
        updateChatList();
      }
    }
  } catch (error) {
    console.error('Error:', error);
    addMessage('Error: Unable to get a response from the server. Please try again.', false);
    isConnected = false;
    updateConnectionStatus('Disconnected', false);
  } finally {
    userInput.disabled = false;
    sendButton.disabled = false;
    userInput.focus();
  }
}

// ============================================
// Upload d'Image
// ============================================
uploadButton.addEventListener('click', () => { imageUpload.click(); });

imageUpload.addEventListener('change', () => {
  const file = imageUpload.files[0];
  if (!file) return;
  const reader = new FileReader();
  reader.onload = (e) => {
    pendingImage = e.target.result;
    imagePreview.innerHTML = `<img src="${pendingImage}" style="max-width:100%; border-radius: var(--border-radius);" />`;
    imagePreview.style.display = "block";
  };
  reader.readAsDataURL(file);
  imageUpload.value = "";
});

// ============================================
// Event Listeners
// ============================================
connectButton.addEventListener('click', () => {
  if (isConnected) {
    isConnected = false;
    updateConnectionStatus('Disconnected', false);
    userInput.disabled = true;
    sendButton.disabled = true;
    addMessage('Disconnected from server.', false);
    currentModel = '';
    modelSelect.disabled = true;
  } else {
    connectToServer();
  }
});

userInput.addEventListener('keypress', (e) => { if (e.key === 'Enter') sendMessage(); });
sendButton.addEventListener('click', sendMessage);
newChatButton.addEventListener('click', () => { createNewChat(); });
toggleSidebarButton.addEventListener('click', () => { chatSidebar.classList.toggle('collapsed'); });

// ============================================
// Initialisation
// ============================================
initSystemPrompt();
setupServerSelection();

serverUrlInput.focus();