src/main.js

import { DEFAULT_LOCAL_MODEL_KEY, LOCAL_MODELS, createPiAgent } from "./piAgent.js";
import { createSandbox } from "./sandbox.js";
import "./styles.css";

const nodes = {
  status: document.querySelector("#status"),
  modelStatus: document.querySelector("#model-status"),
  sandboxStatus: document.querySelector("#sandbox-status"),
  chat: document.querySelector("#chat"),
  eventLog: document.querySelector("#event-log"),
  files: document.querySelector("#files"),
  prompt: document.querySelector("#prompt"),
  composer: document.querySelector("#composer"),
  send: document.querySelector("#send"),
  boot: document.querySelector("#boot-sandbox"),
  reset: document.querySelector("#reset-sandbox"),
  demo: document.querySelector("#demo-prompt"),
  mode: document.querySelector("#mode"),
  device: document.querySelector("#device"),
  gateDevice: document.querySelector("#gate-device"),
  maxTokens: document.querySelector("#max-new-tokens"),
  temperature: document.querySelector("#temperature"),
  modelLabel: document.querySelector("#model-label"),
  loadModel: document.querySelector("#load-model"),
  confirmLoadModel: document.querySelector("#confirm-load-model"),
  useTestModel: document.querySelector("#use-test-model"),
  modelGate: document.querySelector("#model-gate"),
  gateStatus: document.querySelector("#gate-status"),
};

const params = new URLSearchParams(window.location.search);
let modelReady = false;
let transcriptText = "";

if (!navigator.gpu) {
  nodes.device.value = "wasm";
  nodes.gateDevice.value = "wasm";
}
if (params.get("device")) {
  nodes.device.value = params.get("device");
  nodes.gateDevice.value = params.get("device");
}
if (LOCAL_MODELS[params.get("model")]) nodes.mode.value = params.get("model");
if (LOCAL_MODELS[params.get("mode")]) nodes.mode.value = params.get("mode");
if (params.get("mode") === "mock") nodes.mode.value = "mock";
if (!LOCAL_MODELS[nodes.mode.value] && nodes.mode.value !== "mock") nodes.mode.value = DEFAULT_LOCAL_MODEL_KEY;

const sandbox = createSandbox({
  onStatus: (text) => setSandboxStatus(text),
  onLog: (text) => logEvent("sandbox", text),
});

let agent = createAgent();

function selectedModel() {
  return LOCAL_MODELS[nodes.mode.value] || LOCAL_MODELS[DEFAULT_LOCAL_MODEL_KEY];
}

function isMockMode() {
  return nodes.mode.value === "mock";
}

function updateModelLabel() {
  nodes.modelLabel.textContent = isMockMode() ? "Deterministic test model" : selectedModel().id;
}

function textFromContent(content) {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  return content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n");
}

function setStatus(text) {
  nodes.status.textContent = text;
}

function setSandboxStatus(text) {
  nodes.sandboxStatus.textContent = text;
}

function setModelStatus(text) {
  nodes.modelStatus.textContent = text;
}

function setBusy(isBusy) {
  nodes.send.disabled = isBusy;
  nodes.loadModel.disabled = isBusy;
  nodes.confirmLoadModel.disabled = isBusy;
  nodes.useTestModel.disabled = isBusy;
}

function logEvent(kind, text) {
  const line = `[${new Date().toLocaleTimeString()}] ${kind}: ${text}`;
  nodes.eventLog.textContent = `${nodes.eventLog.textContent}${line}\n`;
  nodes.eventLog.scrollTop = nodes.eventLog.scrollHeight;
}

function createAgent() {
  const next = createPiAgent({
    sandbox,
    modelMode: () => nodes.mode.value,
    device: () => nodes.device.value,
    maxTokens: () => nodes.maxTokens.value,
    temperature: () => nodes.temperature.value,
    onModelStatus: setModelStatus,
  });
  next.subscribe((event) => {
    switch (event.type) {
      case "agent_start":
        setStatus("Agent running");
        setBusy(true);
        logEvent("agent", "start");
        break;
      case "message_end":
        renderChat();
        break;
      case "tool_execution_start":
        logEvent("tool", `${event.toolName} started`);
        break;
      case "tool_execution_end":
        logEvent("tool", `${event.toolName} finished`);
        break;
      case "agent_end":
        setStatus("Ready");
        setBusy(false);
        renderChat();
        refreshFiles().catch((error) => logEvent("files", error.message));
        break;
      default:
        break;
    }
  });
  return next;
}

function resetAgent() {
  agent.abort();
  agent = createAgent();
  renderChat();
}

function makeElement(tag, className, text) {
  const element = document.createElement(tag);
  if (className) element.className = className;
  if (text !== undefined) element.textContent = text;
  return element;
}

function renderMessage(message) {
  if (message.role === "user") {
    const bubble = makeElement("article", "message user");
    bubble.append(makeElement("div", "message-label", "You"));
    bubble.append(makeElement("div", "message-body", textFromContent(message.content)));
    return bubble;
  }

  if (message.role === "toolResult") {
    const bubble = makeElement("article", `message tool${message.isError ? " error" : ""}`);
    bubble.append(makeElement("div", "message-label", message.isError ? `Tool error: ${message.toolName}` : `Tool: ${message.toolName}`));
    const body = makeElement("pre", "message-code", textFromContent(message.content));
    bubble.append(body);
    return bubble;
  }

  const bubble = makeElement("article", "message assistant");
  bubble.append(makeElement("div", "message-label", "Pi"));
  const text = textFromContent(message.content);
  if (text) bubble.append(makeElement("div", "message-body", text));
  const toolCalls = Array.isArray(message.content) ? message.content.filter((part) => part.type === "toolCall") : [];
  for (const call of toolCalls) {
    const tool = makeElement("div", "tool-call");
    tool.append(makeElement("span", "tool-name", call.name));
    tool.append(makeElement("code", "", JSON.stringify(call.arguments)));
    bubble.append(tool);
  }
  return bubble;
}

function renderChat() {
  nodes.chat.textContent = "";
  const messages = agent.state.messages;
  transcriptText = messages
    .map((message) => {
      if (message.role === "toolResult") return `TOOL ${message.toolName}\n${textFromContent(message.content)}`;
      return `${message.role.toUpperCase()}\n${textFromContent(message.content)}`;
    })
    .join("\n\n");

  if (messages.length === 0) {
    const empty = makeElement("section", "empty-chat");
    empty.append(makeElement("h2", "", "What should Pi do in the sandbox?"));
    empty.append(makeElement("p", "", "Ready when you are."));
    nodes.chat.append(empty);
  } else {
    for (const message of messages) {
      nodes.chat.append(renderMessage(message));
    }
  }
  nodes.chat.scrollTop = nodes.chat.scrollHeight;
}

async function refreshFiles() {
  if (!sandbox.isReady) {
    nodes.files.textContent = "Sandbox not booted.";
    return;
  }
  nodes.files.textContent = await sandbox.listFiles(".");
}

async function bootSandbox({ refresh = true, rethrow = false } = {}) {
  nodes.boot.disabled = true;
  try {
    await sandbox.boot();
    if (refresh) await refreshFiles();
  } catch (error) {
    setSandboxStatus("Sandbox error");
    logEvent("sandbox", error.stack || error.message || String(error));
    if (rethrow) throw error;
  } finally {
    nodes.boot.disabled = false;
  }
}

function hideGate() {
  nodes.modelGate.classList.add("hidden");
}

function showGate(text = "Ready.") {
  nodes.gateStatus.textContent = text;
  nodes.modelGate.classList.remove("hidden");
}

async function loadModelFromControls() {
  if (nodes.mode.value === "mock") {
    modelReady = true;
    setModelStatus("Deterministic test model");
    hideGate();
    return;
  }

  setBusy(true);
  nodes.gateStatus.textContent = "Downloading model...";
  try {
    nodes.gateStatus.textContent = "Booting sandbox...";
    await bootSandbox({ refresh: false, rethrow: true });
    refreshFiles().catch((error) => logEvent("files", error.message));
    nodes.gateStatus.textContent = "Downloading model...";
    await agent.preloadModel();
    modelReady = true;
    setModelStatus("Model ready");
    nodes.gateStatus.textContent = "Model ready.";
    hideGate();
  } catch (error) {
    const message = error.stack || error.message || String(error);
    setModelStatus("Model error");
    nodes.gateStatus.textContent = message;
    logEvent("model", message);
  } finally {
    setBusy(false);
  }
}

async function sendPrompt() {
  const prompt = nodes.prompt.value.trim();
  if (!prompt) return;
  if (!isMockMode() && !modelReady) {
    showGate("Download the model before sending, or use the test model.");
    return;
  }

  setBusy(true);
  setStatus("Agent running");
  logEvent("agent", "start");
  try {
    await bootSandbox({ refresh: false, rethrow: true });
    refreshFiles().catch((error) => logEvent("files", error.message));
    nodes.prompt.value = "";
    await agent.prompt(prompt);
  } catch (error) {
    setStatus("Error");
    setBusy(false);
    logEvent("agent", error.stack || error.message || String(error));
  }
}

nodes.boot.addEventListener("click", bootSandbox);

nodes.reset.addEventListener("click", async () => {
  nodes.reset.disabled = true;
  try {
    await sandbox.reset();
    resetAgent();
    await refreshFiles();
  } catch (error) {
    logEvent("reset", error.stack || error.message || String(error));
  } finally {
    nodes.reset.disabled = false;
  }
});

nodes.demo.addEventListener("click", () => {
  nodes.prompt.value = "Create hello.js containing JavaScript that computes 21 * 2 and prints result: 42, then run node hello.js.";
  nodes.prompt.focus();
});

nodes.composer.addEventListener("submit", (event) => {
  event.preventDefault();
  setTimeout(() => sendPrompt(), 0);
});

nodes.prompt.addEventListener("keydown", (event) => {
  if (event.key === "Enter" && (event.metaKey || event.ctrlKey)) {
    event.preventDefault();
    setTimeout(() => sendPrompt(), 0);
  }
});

nodes.mode.addEventListener("change", () => {
  modelReady = isMockMode();
  resetAgent();
  updateModelLabel();
  setModelStatus(isMockMode() ? "Deterministic test model" : "Model idle");
});

nodes.device.addEventListener("change", () => {
  nodes.gateDevice.value = nodes.device.value;
  modelReady = isMockMode();
  if (!isMockMode()) setModelStatus("Model idle");
  resetAgent();
});

nodes.gateDevice.addEventListener("change", () => {
  nodes.device.value = nodes.gateDevice.value;
});

nodes.loadModel.addEventListener("click", () => {
  if (isMockMode()) {
    nodes.mode.value = DEFAULT_LOCAL_MODEL_KEY;
    updateModelLabel();
    resetAgent();
  }
  showGate("Ready.");
});

nodes.confirmLoadModel.addEventListener("click", async () => {
  nodes.device.value = nodes.gateDevice.value;
  resetAgent();
  await loadModelFromControls();
});

nodes.useTestModel.addEventListener("click", () => {
  nodes.mode.value = "mock";
  modelReady = true;
  resetAgent();
  setModelStatus("Deterministic test model");
  hideGate();
});

setStatus("Ready");
setSandboxStatus("Not booted");
modelReady = isMockMode();
updateModelLabel();
setModelStatus(modelReady ? "Deterministic test model" : "Model idle");
renderChat();
refreshFiles().catch(() => {});

if (params.get("setup") === "skip" || nodes.mode.value === "mock") {
  hideGate();
} else {
  showGate("Ready.");
}

window.__piWebAgent = {
  get transcript() {
    return transcriptText;
  },
  get modelReady() {
    return modelReady;
  },
};