document.addEventListener("DOMContentLoaded", () => {
  const chatLog = document.getElementById("task1ChatLog");
  const userInput = document.getElementById("task1UserInput");
  const sendBtn = document.getElementById("task1SendBtn");
  const taskSelect = document.getElementById("task1TaskSelect");

  const WELCOME = "Select a task, enter text, and compare the six model responses.";
  let loadingCounter = 0;

  const taskModels = Object.entries(TASK1_MODELS).reduce((acc, [key, model]) => {
    const task = model.task1_task;
    if (!acc[task]) acc[task] = [];
    acc[task].push({ key, ...model });
    return acc;
  }, {});

  const modelOrder = [
    "lr_sentiment", "svm_sentiment", "rf_sentiment", "albert_sentiment", "roberta_sentiment", "distilbert_sentiment",
    "lr_sarcasm", "svm_sarcasm", "rf_sarcasm", "albert_sarcasm", "roberta_sarcasm", "distilbert_sarcasm",
  ];

  Object.keys(taskModels).forEach(task => {
    taskModels[task].sort((a, b) => modelOrder.indexOf(a.key) - modelOrder.indexOf(b.key));
  });

  taskSelect.addEventListener("change", clearChat);

  userInput.addEventListener("input", () => {
    userInput.style.height = "auto";
    userInput.style.height = Math.min(userInput.scrollHeight, 140) + "px";
  });

  userInput.addEventListener("keydown", event => {
    if (event.key === "Enter" && !event.shiftKey) {
      event.preventDefault();
      sendMessage();
    }
  });

  sendBtn.addEventListener("click", sendMessage);

  function clearChat() {
    chatLog.innerHTML = "";
    appendMsg("bot", WELCOME);
  }

  async function sendMessage() {
    const text = userInput.value.trim();
    const task = taskSelect.value;
    const models = taskModels[task] || [];

    if (!text) return;
    if (models.length === 0) {
      appendMsg("bot", "No models are registered for this task.");
      return;
    }

    appendMsg("user", text);
    userInput.value = "";
    userInput.style.height = "auto";
    sendBtn.disabled = true;

    const requests = models.map(model => runModel(model, text, task));
    await Promise.allSettled(requests);

    sendBtn.disabled = false;
    userInput.focus();
  }

  async function runModel(model, text, task) {
    const loadingId = appendLoading(model);

    try {
      const response = await fetch("/api/task1/infer", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ text, task, model_key: model.key }),
      });
      const data = await response.json();
      removeLoading(loadingId);

      if (data.error) {
        appendMsg("bot", `Warning: ${escapeHTML(data.error)}`, model);
      } else {
        appendMsg("bot", formatResult(data, model, task), model);
      }
    } catch (error) {
      removeLoading(loadingId);
      appendMsg("bot", `Warning: Network error: ${escapeHTML(error.message)}`, model);
    }
  }

  function appendMsg(role, content, model = null) {
    const wrapper = document.createElement("div");
    wrapper.className = `msg ${role}`;
    if (model && model.family) wrapper.dataset.family = model.family;

    const avatar = document.createElement("div");
    avatar.className = "msg-avatar";
    avatar.innerHTML = role === "bot"
      ? iconForModel(model)
      : '<i class="fa-solid fa-user"></i>';

    const bubble = document.createElement("div");
    bubble.className = "msg-bubble";
    bubble.innerHTML = role === "bot" ? content : escapeHTML(content);

    wrapper.appendChild(avatar);
    wrapper.appendChild(bubble);
    chatLog.appendChild(wrapper);
    chatLog.scrollTop = chatLog.scrollHeight;
  }

  function appendLoading(model) {
    const id = `task1-loading-${++loadingCounter}`;
    const wrapper = document.createElement("div");
    wrapper.className = "msg bot";
    wrapper.id = id;
    if (model.family) wrapper.dataset.family = model.family;
    wrapper.innerHTML = `
      <div class="msg-avatar">${iconForModel(model)}</div>
      <div class="msg-bubble">
        <div class="model-badge">${escapeHTML(model.name)} <span>${familyLabel(model.family)}</span></div>
        <div class="loading-dots"><span></span><span></span><span></span></div>
      </div>
    `;
    chatLog.appendChild(wrapper);
    chatLog.scrollTop = chatLog.scrollHeight;
    return id;
  }

  function removeLoading(id) {
    const element = document.getElementById(id);
    if (element) element.remove();
  }

  function formatResult(data, model, task) {
    const raw = String(data.label || "").toUpperCase();
    let display = data.label || "Unknown";

    if (task === "sentiment") {
      if (raw === "POSITIVE" || raw === "LABEL_1" || raw === "1") display = "Positive";
      if (raw === "NEGATIVE" || raw === "LABEL_0" || raw === "0") display = "Negative";
    }

    if (task === "sarcasm") {
      if (raw === "SARCASTIC" || raw === "LABEL_1" || raw === "1") display = "Sarcastic";
      if (raw === "NOT_SARCASTIC" || raw === "LABEL_0" || raw === "0") display = "Not Sarcastic";
    }

    const score = data.score !== undefined ? `<span class="rs">Score: ${data.score}</span>` : "";

    return `
      <div class="model-badge">${escapeHTML(model.name)} <span>${familyLabel(model.family)}</span></div>
      <div class="result-block">
        <span class="rl">Result:</span> <span class="rv">${escapeHTML(display)}</span>
        ${score}
      </div>
    `;
  }

  function iconForModel(model) {
    if (!model) return '<i class="fa-solid fa-robot"></i>';
    if (model.family === "classical") return '<i class="fa-solid fa-square-poll-vertical"></i>';
    return '<i class="fa-solid fa-brain"></i>';
  }

  function familyLabel(family) {
    return family === "ptlm" ? "PTLM" : "Classical";
  }

  function escapeHTML(str) {
    return String(str).replace(/[&<>"']/g, tag => ({
      "&": "&amp;",
      "<": "&lt;",
      ">": "&gt;",
      '"': "&quot;",
      "'": "&#39;",
    }[tag]));
  }
});