document.addEventListener("DOMContentLoaded", () => {
  const chatLog = document.getElementById("task2ChatLog");
  const userInput = document.getElementById("task2UserInput");
  const sendBtn = document.getElementById("task2SendBtn");
  const taskSelect = document.getElementById("task2TaskSelect");

  const WELCOME = "Pick a task, send a message, and compare what each variety-trained RoBERTa predicts.";
  let loadingCounter = 0;

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

  const varietyOrder = ["en-AU", "en-IN", "en-UK", "mixed"];
  Object.keys(taskModels).forEach(task => {
    taskModels[task].sort((a, b) =>
      varietyOrder.indexOf(a.trained_on) - varietyOrder.indexOf(b.trained_on)
    );
  });

  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/task2/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}`;

    const avatar = document.createElement("div");
    avatar.className = "msg-avatar";
    avatar.innerHTML = role === "bot"
      ? '<i class="fa-solid fa-brain"></i>'
      : '<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 = `task2-loading-${++loadingCounter}`;
    const wrapper = document.createElement("div");
    wrapper.className = "msg bot";
    wrapper.id = id;
    wrapper.innerHTML = `
      <div class="msg-avatar"><i class="fa-solid fa-brain"></i></div>
      <div class="msg-bubble">
        <div class="model-badge">${escapeHTML(model.name)} <span>Trained on ${escapeHTML(model.trained_on)}</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>` : "";
    const trainedOn = data.trained_on || model.trained_on || "";

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

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