// MedSegAgent demo: ChatGPT-style chat UI driving a real TotalSegmentator run,
// visualized with Niivue. No build step, no framework -- see the implementation
// plan (webapp is a from-scratch additive feature) for context.
//
// Result display pattern: each completed turn shows a small static thumbnail
// inline in its chat bubble (like an image-generation chat) rather than a
// permanently visible viewer panel. Tapping a thumbnail opens a modal with
// the full interactive Niivue viewer (multiplanar view, scrollable slices).
// There is exactly one Niivue instance/canvas for the whole page, reused for
// both off-screen thumbnail rendering and the on-screen modal -- see
// style.css's .viewer-modal (visibility-toggled, never display:none, so the
// canvas keeps a stable non-zero size and never loses its WebGL context).

const chatLog = document.getElementById("chat-log");
const composer = document.getElementById("composer");
const queryInput = document.getElementById("query-input");
const fileInput = document.getElementById("file-input");
const fileChip = document.getElementById("file-chip");
const sendButton = document.getElementById("send-button");
const overlayToggle = document.getElementById("overlay-toggle");
const viewerModal = document.getElementById("viewer-modal");
const viewerModalClose = document.getElementById("viewer-modal-close");

const POLL_INTERVAL_MS = 2000;

// --- Niivue setup -----------------------------------------------------

const nv = new niivue.Niivue({ backColor: [0, 0, 0, 1] });
nv.attachTo("gl");

const SEG_VOLUME_INDEX = 1;

// Identifies whichever turn's volumes are currently loaded into the shared
// canvas, so opening a thumbnail can skip re-fetching the NIfTI files when
// it's already showing that same turn's result.
let loadedJobId = null;

overlayToggle.addEventListener("change", () => {
  if (!nv.volumes[SEG_VOLUME_INDEX]) return;
  nv.volumes[SEG_VOLUME_INDEX].opacity = overlayToggle.checked ? 0.7 : 0;
  nv.updateGLVolume();
});

function openViewerModal() {
  viewerModal.classList.remove("closed");
}

function closeViewerModal() {
  viewerModal.classList.add("closed");
}

viewerModalClose.addEventListener("click", closeViewerModal);
viewerModal.addEventListener("click", (e) => {
  if (e.target === viewerModal) closeViewerModal(); // click on the backdrop
});

// Loads a turn's result into the shared canvas and isolates the requested
// structure (background transparent, target label colored). No-op if this
// job's volumes are already the ones loaded.
async function loadResultIntoViewer(jobId, result) {
  if (loadedJobId === jobId) return;

  await nv.loadVolumes([
    { url: result.input_url, colormap: "gray", opacity: 1 },
    { url: result.seg_url, colormap: "gray", opacity: 0.7 },
  ]);

  // TotalSegmentator was run with roi_subset=[target], so the segmentation
  // volume should only ever contain 0 and this one label id anyway -- this
  // isolation is just to guarantee background is fully transparent.
  const cmap = {
    R: [0, 255],
    G: [0, 70],
    B: [0, 70],
    A: [0, 200],
    I: [0, result.label_id],
    labels: ["background", result.label_name],
  };
  // setColormapLabel is a method on the NVImage instance (nv.volumes[i]),
  // not on the Niivue instance itself -- confirmed against niivue source.
  nv.volumes[SEG_VOLUME_INDEX].setColormapLabel(cmap);
  nv.updateGLVolume();

  overlayToggle.checked = true;
  loadedJobId = jobId;
}

// Renders a turn's result and captures a static PNG snapshot for the chat
// bubble thumbnail. Waits a couple of animation frames after loading so the
// GPU has actually painted before we read the canvas back.
async function renderThumbnail(jobId, result) {
  await loadResultIntoViewer(jobId, result);
  await new Promise((resolve) => requestAnimationFrame(() => requestAnimationFrame(resolve)));
  return nv.canvas.toDataURL("image/png");
}

// --- Chat state ---------------------------------------------------------

let selectedFile = null;
let busy = false;

fileInput.addEventListener("change", () => {
  selectedFile = fileInput.files[0] || null;
  if (selectedFile) {
    fileChip.textContent = selectedFile.name;
    fileChip.hidden = false;
  } else {
    fileChip.hidden = true;
  }
});

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

queryInput.addEventListener("keydown", (e) => {
  if (e.key === "Enter" && !e.shiftKey) {
    e.preventDefault();
    composer.requestSubmit();
  }
});

composer.addEventListener("submit", async (e) => {
  e.preventDefault();
  if (busy) return;

  const query = queryInput.value.trim();
  if (!query) return;
  if (!selectedFile) {
    addBubble("assistant", "先に .nii.gz / .nii ファイルをアップロードしてください。", { error: true });
    return;
  }

  addBubble("user", query, { fileName: selectedFile.name });

  const fileToSend = selectedFile;
  queryInput.value = "";
  queryInput.style.height = "auto";
  fileInput.value = "";
  selectedFile = null;
  fileChip.hidden = true;

  // A new turn is starting: don't leave the modal open showing a stale
  // result while this turn's completion is about to reuse the same canvas.
  closeViewerModal();

  await runChatTurn(query, fileToSend);
});

function addBubble(role, text, opts = {}) {
  const row = document.createElement("div");
  row.className = `bubble-row ${role}`;

  const bubble = document.createElement("div");
  bubble.className = "bubble" + (opts.error ? " error" : "");
  bubble.textContent = text;

  if (opts.fileName) {
    const fileNameEl = document.createElement("span");
    fileNameEl.className = "file-name";
    fileNameEl.textContent = `📎 ${opts.fileName}`;
    bubble.appendChild(fileNameEl);
  }

  row.appendChild(bubble);
  chatLog.appendChild(row);
  chatLog.scrollTop = chatLog.scrollHeight;
  return bubble;
}

function addThumbnailToBubble(bubble, jobId, result, dataUrl) {
  const wrap = document.createElement("div");
  wrap.className = "msg-thumb-wrap";

  const img = document.createElement("img");
  img.className = "msg-thumb";
  img.src = dataUrl;
  img.alt = `${result.label_name} segmentation preview`;
  img.addEventListener("click", async () => {
    await loadResultIntoViewer(jobId, result);
    openViewerModal();
  });
  wrap.appendChild(img);

  const caption = document.createElement("div");
  caption.className = "msg-thumb-caption";
  caption.textContent = "タップして拡大 (断面をスクロール表示)";
  wrap.appendChild(caption);

  bubble.appendChild(wrap);
  chatLog.scrollTop = chatLog.scrollHeight;
}

function setBusy(isBusy) {
  busy = isBusy;
  sendButton.disabled = isBusy;
}

async function runChatTurn(query, file) {
  setBusy(true);
  const statusBubble = addBubble("assistant", "クエリを解析中...");

  const form = new FormData();
  form.append("query", query);
  form.append("file", file, file.name);

  let submitResult;
  try {
    const res = await fetch("/api/chat", { method: "POST", body: form });
    submitResult = await res.json();
    if (!res.ok) {
      throw new Error(submitResult.detail || "リクエストに失敗しました。");
    }
  } catch (err) {
    statusBubble.classList.add("error");
    statusBubble.textContent = `エラー: ${err.message}`;
    setBusy(false);
    return;
  }

  const { job_id, modality, resolved_label } = submitResult;
  statusBubble.textContent =
    `対象を検出しました: ${resolved_label.name} (label ${resolved_label.id}, ${modality})\n` +
    `セグメンテーション実行を待機中...`;

  const startedAt = Date.now();
  await pollJob(job_id, statusBubble, resolved_label, startedAt);
  setBusy(false);
}

function pollJob(jobId, statusBubble, resolvedLabel, startedAt) {
  return new Promise((resolve) => {
    const timer = setInterval(async () => {
      let job;
      try {
        const res = await fetch(`/api/jobs/${jobId}`);
        job = await res.json();
        if (!res.ok) throw new Error(job.detail || "ジョブ情報の取得に失敗しました。");
      } catch (err) {
        clearInterval(timer);
        statusBubble.classList.add("error");
        statusBubble.textContent = `エラー: ${err.message}`;
        resolve();
        return;
      }

      const elapsedSec = Math.round((Date.now() - startedAt) / 1000);

      if (job.status === "queued") {
        statusBubble.textContent = `GPUの空き待ちです... (${elapsedSec}秒)`;
      } else if (job.status === "running") {
        statusBubble.textContent = `セグメンテーション実行中... (${elapsedSec}秒)`;
      } else if (job.status === "done") {
        clearInterval(timer);
        statusBubble.textContent =
          `完了しました（${elapsedSec}秒）: ${resolvedLabel.name} (label ${resolvedLabel.id}) を検出しました。`;
        try {
          const dataUrl = await renderThumbnail(jobId, job.result);
          addThumbnailToBubble(statusBubble, jobId, job.result, dataUrl);
        } catch (err) {
          statusBubble.textContent += `\n(プレビュー画像の生成に失敗しました: ${err.message})`;
        }
        resolve();
      } else if (job.status === "error") {
        clearInterval(timer);
        statusBubble.classList.add("error");
        statusBubble.textContent = `エラー: ${job.message}`;
        resolve();
      }
    }, POLL_INTERVAL_MS);
  });
}