const statusEl = document.getElementById("status");
const messagesEl = document.getElementById("messages");
const formEl = document.getElementById("chat-form");
const inputEl = document.getElementById("chat-input");
const clearEl = document.getElementById("chat-clear");
const sendButtonEl = formEl?.querySelector("button[type='submit']");
const settingsDrawer = document.getElementById("settings-drawer");
const rangeInputs = settingsDrawer
  ? Array.from(settingsDrawer.querySelectorAll("input[type='range']"))
  : [];
const configInputs = {
  topK: document.getElementById("config-top-k"),
  topP: document.getElementById("config-top-p"),
  temperature: document.getElementById("config-temperature"),
  retrieverMaxDocs: document.getElementById("config-retriever-max-docs"),
  rerankerMaxDocs: document.getElementById("config-reranker-max-docs"),
  promptId: document.getElementById("config-prompt"),
};

const scheme = window.location.protocol === "https:" ? "wss" : "ws";
const socketUrl = `${scheme}://${window.location.host}/ws`;
let socket;
let streamingBubble = null;
let loadingBubble = null;
let isAwaitingResponse = false;

const updateSendButtonState = () => {
  if (!sendButtonEl) return;
  sendButtonEl.disabled = isAwaitingResponse;
};

const addMessage = (text, direction) => {
  const bubble = document.createElement("div");
  bubble.className = `message message--${direction}`;
  bubble.textContent = text;
  messagesEl.appendChild(bubble);
  messagesEl.scrollTop = messagesEl.scrollHeight;
};

const showLoadingBubble = () => {
  if (loadingBubble) return;
  loadingBubble = document.createElement("div");
  loadingBubble.className = "message message--in message--loading";
  loadingBubble.innerHTML =
    '<span class="typing-dot"></span><span class="typing-dot"></span><span class="typing-dot"></span>';
  messagesEl.appendChild(loadingBubble);
  messagesEl.scrollTop = messagesEl.scrollHeight;
};

const removeLoadingBubble = () => {
  if (!loadingBubble) return;
  loadingBubble.remove();
  loadingBubble = null;
};

const clearMessages = () => {
  if (!messagesEl) return;
  messagesEl.innerHTML = "";
  streamingBubble = null;
  loadingBubble = null;
};

const connect = () => {
  socket = new WebSocket(socketUrl);

  socket.addEventListener("open", () => {
    statusEl.textContent = "Connected";
  });

  socket.addEventListener("message", (event) => {
    let payload;
    try {
      payload = JSON.parse(event.data);
    } catch (error) {
      addMessage(event.data, "in");
      return;
    }

    if (payload.type === "start") {
      // Keep the loading bubble visible until the first chunk arrives.
      // Just prepare the streaming bubble but don't append it yet.
      streamingBubble = document.createElement("div");
      streamingBubble.className = "message message--in";
      streamingBubble.textContent = "";
      return;
    }

    if (payload.type === "chunk") {
      if (!streamingBubble) {
        streamingBubble = document.createElement("div");
        streamingBubble.className = "message message--in";
      }
      // First chunk: swap loading bubble for the streaming bubble.
      if (!streamingBubble.isConnected) {
        removeLoadingBubble();
        messagesEl.appendChild(streamingBubble);
      }
      streamingBubble.textContent += payload.data ?? "";
      messagesEl.scrollTop = messagesEl.scrollHeight;
      return;
    }

    if (payload.type === "report") {
      console.log(
        "%cREPHRASED QUESTION",
        "border: 2px solid red; padding: 1em; color: red; font-size: 14px; font-weight: bold;",
      );
      console.log(payload.rephrased_question);
      console.log(
        "%cSOURCES",
        "border: 2px solid red; padding: 1em; color: red; font-size: 14px; font-weight: bold;",
      );
      console.log(payload.sources);
      console.log(
        "%cRE-RANKED SOURCES",
        "border: 2px solid red; padding: 1em; color: red; font-size: 14px; font-weight: bold;",
      );
      console.log(payload.reranked_sources);
      return;
    }

    if (payload.type === "end") {
      // Safety net: remove loading bubble if no chunks were ever received.
      removeLoadingBubble();
      streamingBubble = null;
      isAwaitingResponse = false;
      updateSendButtonState();
      return;
    }
  });

  socket.addEventListener("close", () => {
    statusEl.textContent = "Disconnected";
    removeLoadingBubble();
    isAwaitingResponse = false;
    updateSendButtonState();
  });

  socket.addEventListener("error", () => {
    statusEl.textContent = "Connection error";
    removeLoadingBubble();
    isAwaitingResponse = false;
    updateSendButtonState();
  });
};

const updateDrawerValue = (input) => {
  const output = input.parentElement?.querySelector(".drawer__value");
  if (!output) return;
  output.textContent = input.value;
};

rangeInputs.forEach((input) => {
  updateDrawerValue(input);
  input.addEventListener("input", () => updateDrawerValue(input));
});

clearEl?.addEventListener("click", () => {
  clearMessages();
  inputEl?.focus();
  updateSendButtonState();
});

inputEl?.addEventListener("input", updateSendButtonState);

formEl.addEventListener("submit", (event) => {
  event.preventDefault();
  const text = inputEl.value.trim();
  if (!text) return;
  if (!socket || socket.readyState !== WebSocket.OPEN) {
    addMessage("Not connected.", "in");
    return;
  }
  if (isAwaitingResponse) return;
  const config = {
    top_k: Number(configInputs.topK?.value ?? 0),
    top_p: Number(configInputs.topP?.value ?? 0),
    temperature: Number(configInputs.temperature?.value ?? 0),
    retriever_max_docs: Number(configInputs.retrieverMaxDocs?.value ?? 0),
    reranker_max_results: Number(configInputs.rerankerMaxDocs?.value ?? 0),
    prompt_id: Number(configInputs.promptId?.value ?? 0),
  };
  addMessage(text, "out");
  socket.send(JSON.stringify({ message: text, config }));
  inputEl.value = "";
  inputEl.focus();
  isAwaitingResponse = true;
  updateSendButtonState();
  showLoadingBubble();
});

connect();
updateSendButtonState();