const form = document.querySelector("#question-form");
const submitButton = document.querySelector("#submit-button");
const statusChip = document.querySelector("#status-chip");
const statusDetail = document.querySelector("#status-detail");
const authSummary = document.querySelector("#auth-summary");
const loginLink = document.querySelector("#login-link");
const logoutLink = document.querySelector("#logout-link");
const recentSessions = document.querySelector("#recent-sessions");
const answerOutput = document.querySelector("#answer-output");
const sessionIdOutput = document.querySelector("#session-id");
const sentenceList = document.querySelector("#sentence-list");
const traceCount = document.querySelector("#trace-count");
const heatmap = document.querySelector("#heatmap");
const selectionDetails = document.querySelector("#selection-details");
const metrics = document.querySelector("#metrics");
const topEdges = document.querySelector("#top-edges");
const exportJson = document.querySelector("#export-json");
const exportCsv = document.querySelector("#export-csv");

let activeSessionId = null;
let pollHandle = null;
let activePayload = null;
let selectedSentenceIdx = null;
let selectedCell = null;
let currentUser = null;

function currentSentences() {
  return activePayload?.analysis?.sentences || activePayload?.session?.sentences || [];
}

function setStatus(label, detail) {
  statusChip.textContent = label;
  statusDetail.textContent = detail;
}

function clearSelection() {
  selectedSentenceIdx = null;
  selectedCell = null;
  selectionDetails.textContent = "Choose a sentence or heatmap cell.";
  metrics.innerHTML = "";
}

function setExportLinks(sessionId, enabled) {
  exportJson.href = enabled ? `/api/sessions/${sessionId}/export.json` : "#";
  exportCsv.href = enabled ? `/api/sessions/${sessionId}/export.csv` : "#";
  exportJson.classList.toggle("is-disabled", !enabled);
  exportCsv.classList.toggle("is-disabled", !enabled);
  exportJson.setAttribute("aria-disabled", String(!enabled));
  exportCsv.setAttribute("aria-disabled", String(!enabled));
}

function renderRecentSessions(sessions) {
  recentSessions.innerHTML = "";
  if (!currentUser?.authenticated) {
    recentSessions.textContent = currentUser?.auth_required
      ? "Sign in to view your jobs."
      : "Anonymous mode is enabled.";
    return;
  }
  if (!sessions.length) {
    recentSessions.textContent = "No jobs yet on this running Space instance.";
    return;
  }
  sessions.forEach((session) => {
    const item = document.createElement("button");
    item.type = "button";
    item.className = "session-card";
    if (session.id === activeSessionId) {
      item.classList.add("is-active");
    }
    item.innerHTML = `<strong>${session.status}</strong>${session.question}`;
    item.addEventListener("click", () => {
      startPolling(session.id);
    });
    recentSessions.appendChild(item);
  });
}

async function loadRecentSessions() {
  if (!currentUser?.authenticated && currentUser?.auth_required) {
    renderRecentSessions([]);
    return;
  }
  const response = await fetch("/api/sessions?limit=8");
  if (!response.ok) {
    renderRecentSessions([]);
    return;
  }
  renderRecentSessions(await response.json());
}

function renderAuth() {
  if (!currentUser) {
    authSummary.textContent = "Checking sign-in status.";
    return;
  }
  loginLink.hidden = currentUser.authenticated;
  logoutLink.hidden = !currentUser.authenticated;
  if (currentUser.authenticated) {
    authSummary.textContent = `Signed in as ${currentUser.full_name || currentUser.username}.`;
    return;
  }
  authSummary.textContent = currentUser.auth_required
    ? "Sign in with Hugging Face to run analysis jobs."
    : "Anonymous access is enabled for this Space.";
}

function renderSession(session) {
  sessionIdOutput.textContent = session.id ? `Session ${session.id.slice(0, 8)}` : "";
  answerOutput.textContent = session.answer || "Waiting for generated answer.";
  const sentences = currentSentences();
  traceCount.textContent = sentences.length ? `${sentences.length} sentences` : "";

  sentenceList.innerHTML = "";
  sentences.forEach((sentence, index) => {
    const item = document.createElement("li");
    item.className = "sentence-item";
    item.innerHTML = `<strong>${index}</strong> ${sentence}`;
    item.addEventListener("click", () => {
      selectedSentenceIdx = index;
      selectedCell = null;
      renderSelection();
    });
    sentenceList.appendChild(item);
  });
}

function colorForValue(value, maxAbs) {
  if (!maxAbs) {
    return "rgba(224, 223, 218, 0.8)";
  }
  const normalized = Math.max(-1, Math.min(1, value / maxAbs));
  if (normalized >= 0) {
    const alpha = 0.12 + normalized * 0.88;
    return `rgba(14, 90, 138, ${alpha.toFixed(3)})`;
  }
  const alpha = 0.12 + Math.abs(normalized) * 0.88;
  return `rgba(215, 106, 52, ${alpha.toFixed(3)})`;
}

function renderHeatmap(result) {
  const matrix = result?.suppression_matrix;
  if (!matrix || !matrix.length) {
    heatmap.className = "heatmap placeholder-box";
    heatmap.textContent = "Analysis pending.";
    return;
  }

  const flatValues = matrix.flat();
  const maxAbs = Math.max(...flatValues.map((value) => Math.abs(value)));
  heatmap.className = "heatmap";
  heatmap.innerHTML = "";
  const grid = document.createElement("div");
  grid.className = "heatmap-grid";
  grid.style.gridTemplateColumns = `repeat(${matrix.length}, 32px)`;

  matrix.forEach((row, rowIndex) => {
    row.forEach((value, colIndex) => {
      const cell = document.createElement("button");
      cell.type = "button";
      cell.className = "heatmap-cell";
      cell.style.background = colorForValue(value, maxAbs);
      cell.title = `target ${rowIndex} ← source ${colIndex}: ${value.toFixed(4)}`;
      if (selectedCell && selectedCell.row === rowIndex && selectedCell.col === colIndex) {
        cell.classList.add("is-selected");
      }
      cell.addEventListener("click", () => {
        selectedSentenceIdx = null;
        selectedCell = { row: rowIndex, col: colIndex, value };
        renderSelection();
      });
      grid.appendChild(cell);
    });
  });
  heatmap.appendChild(grid);
}

function renderTopEdges(result) {
  topEdges.innerHTML = "";
  const edges = result?.top_edges || [];
  if (!edges.length) {
    return;
  }
  edges.slice(0, 5).forEach((edge) => {
    const item = document.createElement("div");
    item.className = "edge-card";
    item.innerHTML = `<strong>${edge.source_sentence_idx} → ${edge.target_sentence_idx}</strong>${edge.score.toFixed(4)}`;
    topEdges.appendChild(item);
  });
}

function renderSelection() {
  Array.from(sentenceList.children).forEach((item, index) => {
    item.classList.toggle("is-active", selectedSentenceIdx === index);
  });

  const result = activePayload?.analysis;
  const session = activePayload?.session;
  if (!session) {
    clearSelection();
    return;
  }

  metrics.innerHTML = "";
  if (selectedSentenceIdx != null && result) {
    const outgoing = result.outgoing_importance[selectedSentenceIdx] ?? 0;
    const incoming = result.incoming_importance[selectedSentenceIdx] ?? 0;
    selectionDetails.innerHTML = `<strong>Sentence ${selectedSentenceIdx}</strong><br>${currentSentences()[selectedSentenceIdx] || ""}`;
    metrics.innerHTML = `
      <div class="metric-card"><strong>Outgoing impact</strong>${outgoing.toFixed(4)}</div>
      <div class="metric-card"><strong>Incoming dependence</strong>${incoming.toFixed(4)}</div>
    `;
    return;
  }

  if (selectedCell) {
    selectionDetails.innerHTML = `<strong>Edge ${selectedCell.col} → ${selectedCell.row}</strong><br>Influence score ${selectedCell.value.toFixed(4)}`;
    metrics.innerHTML = `
      <div class="metric-card"><strong>Source sentence</strong>${selectedCell.col}</div>
      <div class="metric-card"><strong>Target sentence</strong>${selectedCell.row}</div>
    `;
    return;
  }

  selectionDetails.textContent = "Choose a sentence or heatmap cell.";
}

async function fetchSession(sessionId) {
  const response = await fetch(`/api/sessions/${sessionId}/result`);
  if (!response.ok) {
    throw new Error(`Failed to fetch session ${sessionId}`);
  }
  return response.json();
}

function updateFromPayload(payload) {
  activePayload = payload;
  activeSessionId = payload.session.id;
  renderSession(payload.session);
  renderHeatmap(payload.analysis);
  renderTopEdges(payload.analysis);
  renderSelection();
  setExportLinks(payload.session.id, payload.session.status === "completed");

  const { status, error } = payload.session;
  if (status === "queued") setStatus("Queued", "Waiting for a worker slot.");
  if (status === "generating") setStatus("Generating", "The model is producing an answer and visible trace.");
  if (status === "answer_ready") setStatus("Analysis pending", "Answer is ready. Attribution analysis is starting.");
  if (status === "analyzing") setStatus("Analyzing", "Running forward and backward passes for sentence influence.");
  if (status === "completed") setStatus("Completed", "Analysis finished.");
  if (status === "failed") setStatus("Failed", error || "The session failed.");

  if (["completed", "failed"].includes(status) && pollHandle) {
    window.clearInterval(pollHandle);
    pollHandle = null;
  }

  loadRecentSessions().catch(() => {});
}

async function startPolling(sessionId) {
  activeSessionId = sessionId;
  if (pollHandle) {
    window.clearInterval(pollHandle);
  }
  const tick = async () => {
    try {
      const payload = await fetchSession(sessionId);
      updateFromPayload(payload);
    } catch (error) {
      setStatus("Error", String(error));
      window.clearInterval(pollHandle);
      pollHandle = null;
    }
  };
  await tick();
  pollHandle = window.setInterval(tick, 2500);
}

form.addEventListener("submit", async (event) => {
  event.preventDefault();
  if (!currentUser?.authenticated && currentUser?.auth_required) {
    window.location.href = loginLink.href;
    return;
  }
  submitButton.disabled = true;
  clearSelection();
  sentenceList.innerHTML = "";
  heatmap.className = "heatmap placeholder-box";
  heatmap.textContent = "Analysis pending.";
  topEdges.innerHTML = "";
  answerOutput.textContent = "Waiting for generated answer.";
  sessionIdOutput.textContent = "";
  traceCount.textContent = "";
  setExportLinks("", false);
  setStatus("Submitting", "Creating session.");

  const payload = {
    question: document.querySelector("#question").value,
    max_new_tokens: Number(document.querySelector("#max-new-tokens").value),
    max_trace_tokens: Number(document.querySelector("#max-trace-tokens").value),
    max_sentences: Number(document.querySelector("#max-sentences").value),
  };

  try {
    const response = await fetch("/api/sessions", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify(payload),
    });
    if (!response.ok) {
      throw new Error(await response.text());
    }
    const session = await response.json();
    await startPolling(session.id);
  } catch (error) {
    setStatus("Error", String(error));
  } finally {
    submitButton.disabled = false;
  }
});

async function initialize() {
  setExportLinks("", false);
  try {
    const response = await fetch("/api/me");
    currentUser = await response.json();
  } catch (_error) {
    currentUser = {
      authenticated: false,
      auth_required: true,
      login_url: "/oauth/huggingface/login",
      logout_url: "/oauth/huggingface/logout",
    };
  }
  renderAuth();
  await loadRecentSessions();
}

initialize();