import {
  env,
  pipeline,
  TextStreamer,
} from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0";

const MODEL_ID = "onnx-community/gemma-3-270m-it-ONNX";
const MODEL_REVISION = "2dbbfdb1b59bd034eb959428c6a7da9dd7ea27f0";
const STORE_KEY = "spinor-gpt-sessions-v1";
const SETTINGS_KEY = "spinor-gpt-settings-v1";
const MAX_HISTORY_MESSAGES = 12;
const MAX_EVIDENCE_CHARS = 9000;

env.allowRemoteModels = true;
env.allowLocalModels = false;
env.useBrowserCache = true;

const $ = (selector) => document.querySelector(selector);
const elements = {
  sidebar: $("#sidebar"),
  history: $("#history"),
  newChat: $("#new-chat"),
  exportChat: $("#export-chat"),
  clearData: $("#clear-data"),
  mobileMenu: $("#mobile-menu"),
  mode: $("#runtime-mode"),
  modelTitle: $("#model-title"),
  status: $("#runtime-status"),
  modeFootnote: $("#mode-footnote"),
  empty: $("#empty-state"),
  messages: $("#messages"),
  stage: $("#chat-stage"),
  prompt: $("#prompt"),
  send: $("#send"),
  files: $("#knowledge-files"),
  knowledgeStrip: $("#knowledge-strip"),
  settingsButton: $("#settings-button"),
  closeSettings: $("#close-settings"),
  settingsPanel: $("#settings-panel"),
  scrim: $("#scrim"),
  system: $("#system-prompt"),
  temperature: $("#temperature"),
  temperatureValue: $("#temperature-value"),
  maxTokens: $("#max-tokens"),
  maxTokensValue: $("#max-tokens-value"),
  nativeSettings: $("#native-settings"),
  nativeEndpoint: $("#native-endpoint"),
  evidenceLock: $("#evidence-lock"),
  settingsRuntimeName: $("#settings-runtime-name"),
  settingsRuntimeDetail: $("#settings-runtime-detail"),
  loadModal: $("#load-modal"),
  loadMessage: $("#load-message"),
  progressFill: $("#progress-fill"),
  progressDetail: $("#progress-detail"),
  toast: $("#toast"),
};

const state = {
  generator: null,
  loadingPromise: null,
  generating: false,
  sessions: [],
  currentId: null,
  evidence: [],
  fileProgress: new Map(),
};

function createId() {
  return `${Date.now().toString(36)}-${crypto.getRandomValues(new Uint32Array(1))[0].toString(36)}`;
}

function defaultSession() {
  return { id: createId(), title: "New conversation", messages: [], updated: Date.now() };
}

function currentSession() {
  return state.sessions.find((session) => session.id === state.currentId);
}

function loadState() {
  try {
    state.sessions = JSON.parse(localStorage.getItem(STORE_KEY)) || [];
  } catch {
    state.sessions = [];
  }
  state.sessions = state.sessions.map((session) => ({
    ...session,
    messages: (session.messages || []).filter(
      (message) => !(message.role === "assistant" && /^Runtime error:/i.test(message.content || "")),
    ),
  }));
  if (!state.sessions.length) state.sessions = [defaultSession()];
  state.currentId = state.sessions[0].id;
  try {
    const settings = JSON.parse(localStorage.getItem(SETTINGS_KEY)) || {};
    elements.mode.value = settings.mode || "auto";
    if (settings.system) elements.system.value = settings.system;
    if (settings.temperature != null) elements.temperature.value = settings.temperature;
    if (settings.maxTokens != null) elements.maxTokens.value = settings.maxTokens;
    if (settings.endpoint) elements.nativeEndpoint.value = settings.endpoint;
    elements.evidenceLock.checked = Boolean(settings.evidenceLock);
    if (
      location.protocol === "https:"
      && elements.mode.value === "native"
      && elements.nativeEndpoint.value.trim().startsWith("http://")
    ) {
      elements.mode.value = "auto";
    }
  } catch {}
}

function persist() {
  localStorage.setItem(STORE_KEY, JSON.stringify(state.sessions));
  localStorage.setItem(
    SETTINGS_KEY,
    JSON.stringify({
      mode: elements.mode.value,
      system: elements.system.value,
      temperature: Number(elements.temperature.value),
      maxTokens: Number(elements.maxTokens.value),
      endpoint: elements.nativeEndpoint.value.trim(),
      evidenceLock: elements.evidenceLock.checked,
    }),
  );
}

function escapeHtml(value) {
  return value
    .replaceAll("&", "&")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;");
}

function runtimeErrorMessage(error) {
  if (typeof error === "string" && error.trim()) return error;
  if (error?.message) return String(error.message);
  if (error?.name) return String(error.name);
  return "The browser model could not start. Reload the page and try Browser WebGPU again.";
}

function renderMarkdown(value) {
  const blocks = escapeHtml(value).split(/```/);
  return blocks
    .map((block, index) => {
      if (index % 2) return `<pre><code>${block.replace(/^\w+\n/, "")}</code></pre>`;
      return block
        .split(/\n{2,}/)
        .map((paragraph) => `<p>${paragraph
          .replace(/`([^`]+)`/g, "<code>$1</code>")
          .replace(/\*\*([^*]+)\*\*/g, "<strong>$1</strong>")
          .replaceAll("\n", "<br>")}</p>`)
        .join("");
    })
    .join("");
}

function renderSessions() {
  const ordered = [...state.sessions].sort((a, b) => b.updated - a.updated);
  elements.history.replaceChildren(
    ...ordered.map((session) => {
      const button = document.createElement("button");
      button.className = `history-item${session.id === state.currentId ? " active" : ""}`;
      button.textContent = session.title;
      button.onclick = () => {
        state.currentId = session.id;
        state.evidence = [];
        renderAll();
        elements.sidebar.classList.remove("open");
      };
      return button;
    }),
  );
}

function makeMessage(message, streaming = false) {
  const article = document.createElement("article");
  article.className = `message ${message.role}`;
  const avatar = document.createElement("div");
  avatar.className = "avatar";
  avatar.textContent = message.role === "assistant" ? "ψ" : "You";
  const body = document.createElement("div");
  const head = document.createElement("div");
  head.className = "message-head";
  head.textContent = message.role === "assistant" ? "Spinor GPT" : "You";
  const content = document.createElement("div");
  content.className = `message-content${streaming ? " cursor" : ""}`;
  content.innerHTML = renderMarkdown(message.content || "");
  body.append(head, content);
  article.append(avatar, body);
  return { article, content };
}

function renderMessages() {
  const session = currentSession();
  const hasMessages = Boolean(session?.messages.length);
  elements.empty.classList.toggle("hidden", hasMessages);
  elements.messages.replaceChildren();
  if (hasMessages) {
    for (const message of session.messages) elements.messages.append(makeMessage(message).article);
    requestAnimationFrame(() => { elements.stage.scrollTop = elements.stage.scrollHeight; });
  }
}

function renderEvidence() {
  elements.knowledgeStrip.classList.toggle("hidden", !state.evidence.length);
  elements.knowledgeStrip.replaceChildren(
    ...state.evidence.map((file) => {
      const chip = document.createElement("span");
      chip.className = "knowledge-chip";
      chip.textContent = `${file.name} · ${Math.ceil(file.text.length / 1000)}k`;
      return chip;
    }),
  );
}

function renderAll() {
  renderSessions();
  renderMessages();
  renderEvidence();
  updateModeUI();
}

function setStatus(label, kind = "") {
  elements.status.className = `status-pill ${kind}`;
  elements.status.innerHTML = `<i></i>${escapeHtml(label)}`;
}

function toast(message) {
  elements.toast.textContent = message;
  elements.toast.classList.remove("hidden");
  clearTimeout(toast.timeout);
  toast.timeout = setTimeout(() => elements.toast.classList.add("hidden"), 2600);
}

function nativeEndpointBlocked() {
  return location.protocol === "https:"
    && elements.nativeEndpoint.value.trim().startsWith("http://");
}

async function useBrowserFallback(messages, contentElement, reason) {
  elements.mode.value = "browser";
  updateModeUI();
  toast(reason);
  return browserGenerate(messages, contentElement);
}

function updateModeUI() {
  if (elements.mode.value === "native" && nativeEndpointBlocked()) {
    elements.mode.value = "browser";
    toast("Native Spinor needs a local page or HTTPS endpoint. Using Browser WebGPU here.");
  }
  const native = elements.mode.value === "native";
  const automatic = elements.mode.value === "auto";
  elements.modelTitle.textContent = native ? "Native Spinor · MLX" : automatic ? "Auto relay" : "Browser WebGPU";
  elements.modeFootnote.textContent = native
    ? "Connected only to your configured endpoint"
    : automatic ? "Native when reachable · browser fallback when offline" : "Browser model runs on your device";
  elements.nativeSettings.classList.toggle("hidden", !native && !automatic);
  elements.settingsRuntimeName.textContent = native ? "Spinor GPT 7B · MLX q4" : automatic ? "Native + browser failover" : "Gemma 3 270M · FP16";
  elements.settingsRuntimeDetail.textContent = native
    ? "Apple MLX · quaternion query rotation"
    : automatic ? "HTTPS relay · FP16 WebGPU fallback" : `${navigator.gpu ? "WebGPU" : "WebGPU unavailable"} · Transformers.js`;
  setStatus(native ? "Endpoint idle" : automatic ? "Failover armed" : state.generator ? "Model loaded" : "Ready to load", state.generator && !native ? "ready" : "");
  persist();
}

function updateLoadingProgress(info) {
  if (info.status === "progress") {
    state.fileProgress.set(info.file, Number(info.progress) || 0);
    const values = [...state.fileProgress.values()];
    const average = values.reduce((sum, value) => sum + value, 0) / values.length;
    elements.progressFill.style.width = `${Math.max(3, average)}%`;
    elements.loadMessage.textContent = `Downloading ${info.file.split("/").pop()}…`;
    if (info.loaded && info.total) {
      elements.progressDetail.textContent = `${(info.loaded / 1048576).toFixed(1)} of ${(info.total / 1048576).toFixed(1)} MB · cached after first use`;
    }
  } else if (info.status === "ready") {
    elements.progressFill.style.width = "100%";
    elements.loadMessage.textContent = "Model ready";
  }
}

async function ensureBrowserModel() {
  if (state.generator) return state.generator;
  if (state.loadingPromise) return state.loadingPromise;
  if (!navigator.gpu) {
    throw new Error("Browser inference requires WebGPU. Open this Space in a current Chrome, Edge, or Safari browser, or connect a native HTTPS endpoint.");
  }
  elements.loadModal.classList.remove("hidden");
  elements.scrim.classList.remove("hidden");
  state.fileProgress.clear();
  setStatus("Loading model");
  state.loadingPromise = pipeline("text-generation", MODEL_ID, {
    device: "webgpu",
    dtype: "fp16",
    revision: MODEL_REVISION,
    progress_callback: updateLoadingProgress,
  })
    .then((generator) => {
      state.generator = generator;
      setStatus("Model loaded", "ready");
      return generator;
    })
    .finally(() => {
      state.loadingPromise = null;
      setTimeout(() => {
        elements.loadModal.classList.add("hidden");
        if (!elements.settingsPanel.classList.contains("open")) elements.scrim.classList.add("hidden");
      }, 350);
    });
  return state.loadingPromise;
}

function evidenceSystemPrompt() {
  if (!state.evidence.length) return "";
  const combined = state.evidence
    .map((file, index) => `[S${index + 1}] ${file.name}\n${file.text}`)
    .join("\n\n")
    .slice(0, MAX_EVIDENCE_CHARS);
  const rule = elements.evidenceLock.checked
    ? "Use only this evidence for factual claims. Cite [S1], [S2], etc. If it does not answer the request, say: Insufficient evidence."
    : "Use the attached material when relevant and cite it as [S1], [S2], etc.";
  return `\n\n${rule}\n\nAttached evidence:\n${combined}`;
}

async function browserGenerate(messages, contentElement) {
  const generator = await ensureBrowserModel();
  let generated = "";
  const streamer = new TextStreamer(generator.tokenizer, {
    skip_prompt: true,
    skip_special_tokens: true,
    callback_function: (token) => {
      generated += token;
      contentElement.innerHTML = renderMarkdown(generated);
      elements.stage.scrollTop = elements.stage.scrollHeight;
    },
  });
  const temperature = Number(elements.temperature.value);
  const result = await generator(messages, {
    max_new_tokens: Number(elements.maxTokens.value),
    do_sample: temperature > 0,
    temperature: Math.max(temperature, 0.01),
    top_p: 0.9,
    repetition_penalty: 1.08,
    streamer,
  });
  if (!generated) {
    const output = result?.[0]?.generated_text;
    if (Array.isArray(output)) generated = output.at(-1)?.content || "";
    else generated = String(output || "");
  }
  return generated.trim();
}

async function nativeGenerate(messages, contentElement) {
  const endpoint = elements.nativeEndpoint.value.trim().replace(/\/$/, "");
  if (nativeEndpointBlocked()) {
    return useBrowserFallback(
      messages,
      contentElement,
      "Native endpoint is unavailable from this hosted page — using Browser WebGPU.",
    );
  }
  setStatus("Generating");
  let response;
  try {
    response = await fetch(`${endpoint}/api/chat`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        model: "spinor-gpt:7b",
        messages,
        stream: false,
        temperature: Number(elements.temperature.value),
        num_predict: Number(elements.maxTokens.value),
      }),
    });
  } catch (error) {
    console.warn("Native Spinor endpoint is offline", error);
    return useBrowserFallback(
      messages,
      contentElement,
      "Native Spinor is offline — using Browser WebGPU.",
    );
  }
  if (!response.ok) throw new Error(`Native endpoint returned HTTP ${response.status}`);
  const payload = await response.json();
  if (payload.error) throw new Error(payload.error);
  setStatus(`${Number(payload.tokens_per_second || 0).toFixed(1)} tok/s`, "ready");
  return String(payload.message?.content || payload.response || "").trim();
}

async function nativeReachable() {
  const endpoint = elements.nativeEndpoint.value.trim().replace(/\/$/, "");
  if (!endpoint) return false;
  if (location.protocol === "https:" && endpoint.startsWith("http://")) return false;
  const controller = new AbortController();
  const timeout = setTimeout(() => controller.abort(), 1600);
  try {
    const response = await fetch(`${endpoint}/health`, { signal: controller.signal });
    return response.ok;
  } catch {
    return false;
  } finally {
    clearTimeout(timeout);
  }
}

async function autoGenerate(messages, contentElement) {
  setStatus("Checking native relay");
  if (await nativeReachable()) {
    setStatus("Native relay online", "ready");
    return nativeGenerate(messages, contentElement);
  }
  toast("Native relay is offline — using the private browser model.");
  setStatus("Browser fallback");
  return browserGenerate(messages, contentElement);
}

async function sendMessage() {
  const prompt = elements.prompt.value.trim();
  if (!prompt || state.generating) return;
  const session = currentSession();
  if (!session.messages.length) session.title = prompt.slice(0, 38) + (prompt.length > 38 ? "…" : "");
  session.messages.push({ role: "user", content: prompt });
  session.updated = Date.now();
  elements.prompt.value = "";
  autoResize();
  state.generating = true;
  elements.send.disabled = true;
  renderAll();

  const assistant = { role: "assistant", content: "" };
  const rendered = makeMessage(assistant, true);
  elements.messages.append(rendered.article);
  elements.stage.scrollTop = elements.stage.scrollHeight;

  const system = elements.system.value.trim() + evidenceSystemPrompt();
  const usableHistory = session.messages
    .filter((message) => !(message.role === "assistant" && /^Runtime error:/i.test(message.content || "")))
    .slice(-MAX_HISTORY_MESSAGES);
  const conversation = [
    { role: "system", content: system },
    ...usableHistory,
  ];
  try {
    const answer = elements.mode.value === "browser"
      ? await browserGenerate(conversation, rendered.content)
      : elements.mode.value === "native"
        ? await nativeGenerate(conversation, rendered.content)
        : await autoGenerate(conversation, rendered.content);
    assistant.content = answer || "I could not produce a response.";
    session.messages.push(assistant);
    rendered.content.innerHTML = renderMarkdown(assistant.content);
    rendered.content.classList.remove("cursor");
  } catch (error) {
    const message = runtimeErrorMessage(error);
    console.error("Spinor GPT runtime failure", error);
    assistant.content = `Runtime error: ${message}`;
    session.messages.push(assistant);
    rendered.content.innerHTML = renderMarkdown(assistant.content);
    rendered.content.classList.remove("cursor");
    setStatus("Runtime error", "error");
    toast(message);
  } finally {
    state.generating = false;
    elements.send.disabled = false;
    session.updated = Date.now();
    persist();
    renderSessions();
    elements.stage.scrollTop = elements.stage.scrollHeight;
  }
}

function autoResize() {
  elements.prompt.style.height = "auto";
  elements.prompt.style.height = `${Math.min(elements.prompt.scrollHeight, 160)}px`;
}

function newConversation() {
  const session = defaultSession();
  state.sessions.unshift(session);
  state.currentId = session.id;
  state.evidence = [];
  persist();
  renderAll();
  elements.prompt.focus();
}

function openSettings(open) {
  elements.settingsPanel.classList.toggle("open", open);
  elements.scrim.classList.toggle("hidden", !open);
}

async function attachFiles(event) {
  const files = [...event.target.files];
  const loaded = [];
  for (const file of files) {
    if (file.size > 1_000_000) {
      toast(`${file.name} is larger than 1 MB and was skipped.`);
      continue;
    }
    loaded.push({ name: file.name, text: (await file.text()).slice(0, MAX_EVIDENCE_CHARS) });
  }
  state.evidence = [...state.evidence, ...loaded].slice(-5);
  renderEvidence();
  event.target.value = "";
}

function exportConversation() {
  const session = currentSession();
  const text = session.messages.map((message) => `## ${message.role === "user" ? "You" : "Spinor GPT"}\n\n${message.content}`).join("\n\n");
  const blob = new Blob([`# ${session.title}\n\n${text}\n`], { type: "text/markdown" });
  const link = document.createElement("a");
  link.href = URL.createObjectURL(blob);
  link.download = `${session.title.replace(/[^a-z0-9]+/gi, "-").toLowerCase() || "spinor-chat"}.md`;
  link.click();
  URL.revokeObjectURL(link.href);
}

elements.newChat.onclick = newConversation;
elements.send.onclick = sendMessage;
elements.prompt.addEventListener("input", autoResize);
elements.prompt.addEventListener("keydown", (event) => {
  if (event.key === "Enter" && !event.shiftKey) {
    event.preventDefault();
    sendMessage();
  }
});
elements.files.onchange = attachFiles;
elements.mode.onchange = updateModeUI;
elements.settingsButton.onclick = () => openSettings(true);
elements.closeSettings.onclick = () => openSettings(false);
elements.scrim.onclick = () => openSettings(false);
elements.mobileMenu.onclick = () => elements.sidebar.classList.toggle("open");
elements.exportChat.onclick = exportConversation;
elements.clearData.onclick = () => {
  if (!confirm("Delete every locally saved conversation and setting?")) return;
  localStorage.removeItem(STORE_KEY);
  localStorage.removeItem(SETTINGS_KEY);
  state.sessions = [defaultSession()];
  state.currentId = state.sessions[0].id;
  state.evidence = [];
  renderAll();
};
elements.temperature.oninput = () => { elements.temperatureValue.value = elements.temperature.value; persist(); };
elements.maxTokens.oninput = () => { elements.maxTokensValue.value = elements.maxTokens.value; persist(); };
elements.system.onchange = persist;
elements.nativeEndpoint.onchange = updateModeUI;
elements.evidenceLock.onchange = persist;
document.querySelectorAll(".suggestion").forEach((button) => {
  button.onclick = () => {
    elements.prompt.value = button.dataset.prompt;
    autoResize();
    elements.prompt.focus();
  };
});
window.addEventListener("pagehide", () => {
  if (state.generator) state.generator.dispose();
});

loadState();
elements.temperatureValue.value = elements.temperature.value;
elements.maxTokensValue.value = elements.maxTokens.value;
renderAll();
elements.prompt.focus();