src/main.ts

import "./styles.css";
import {
  DEFAULT_PROMPT,
  DEFAULT_SETTINGS,
  DEFAULT_SYSTEM_PROMPT,
  MODEL_ID,
  PRESETS,
  type GenerationSettings,
  type Precision,
} from "./model";
import type { WorkerRequest, WorkerResponse } from "./messages";

const app = document.querySelector<HTMLDivElement>("#app");

if (!app) {
  throw new Error("Missing #app root");
}

let jobId = 0;
let isGenerating = false;
let generatedText = "";
const diagnostics: string[] = [];
const worker = new Worker(new URL("./worker.ts", import.meta.url), { type: "module" });

app.innerHTML = `
  <main class="shell">
    <section class="console" aria-labelledby="app-title">
      <header class="topbar">
        <div>
          <p class="eyebrow">Transformers.js / WebGPU</p>
          <h1 id="app-title">LFM2.5 230M</h1>
        </div>
        <div class="runtime" id="runtime-status">
          <span class="runtime-dot" aria-hidden="true"></span>
          <span id="runtime-label">Checking WebGPU</span>
        </div>
      </header>

      <div class="workbench">
        <form class="controls" id="generation-form">
          <label class="field field-block">
            <span>System</span>
            <textarea id="system-prompt" rows="3" spellcheck="true"></textarea>
          </label>

          <div class="preset-row" role="group" aria-label="Prompt presets">
            ${PRESETS.map(
              (preset) =>
                `<button class="preset" type="button" data-preset="${preset.id}">${preset.label}</button>`,
            ).join("")}
          </div>

          <label class="field field-block prompt-field">
            <span>Prompt</span>
            <textarea id="prompt" rows="9" spellcheck="true"></textarea>
          </label>

          <div class="setting-grid">
            <label class="field">
              <span>Precision</span>
              <select id="precision">
                <option value="q4">Q4</option>
                <option value="fp16">FP16</option>
              </select>
            </label>

            <label class="field">
              <span>Max tokens</span>
              <input id="max-new-tokens" type="number" min="16" max="1024" step="8" />
            </label>

            <label class="field">
              <span>Temperature</span>
              <input id="temperature" type="number" min="0" max="1.5" step="0.05" />
            </label>

            <label class="field">
              <span>Top K</span>
              <input id="top-k" type="number" min="1" max="200" step="1" />
            </label>

            <label class="field">
              <span>Penalty</span>
              <input id="repetition-penalty" type="number" min="1" max="2" step="0.01" />
            </label>
          </div>

          <div class="actions">
            <button class="primary" id="generate" type="submit">Generate</button>
            <button class="secondary" id="stop" type="button">Stop</button>
            <button class="secondary" id="clear" type="button">Clear</button>
          </div>
        </form>

        <section class="output-panel" aria-label="Generation output">
          <div class="output-header">
            <div>
              <p class="eyebrow">Local output</p>
              <h2>Response</h2>
            </div>
            <div class="metric" id="metric">Idle</div>
          </div>
          <div class="progress-shell" aria-hidden="true">
            <div class="progress-bar" id="progress-bar"></div>
          </div>
          <pre class="output" id="output" tabindex="0"></pre>
          <div class="token-rail" id="token-rail" aria-hidden="true"></div>
          <section class="diagnostics-panel" aria-label="Diagnostics">
            <div class="diagnostics-header">
              <div>
                <p class="eyebrow">Diagnostics</p>
                <h2>Debug log</h2>
              </div>
              <button class="secondary" id="copy-diagnostics" type="button">Copy log</button>
            </div>
            <pre class="diagnostics" id="diagnostics" tabindex="0"></pre>
          </section>
        </section>
      </div>
    </section>
  </main>
`;

const runtimeStatus = element<HTMLDivElement>("#runtime-status");
const runtimeLabel = element<HTMLSpanElement>("#runtime-label");
const form = element<HTMLFormElement>("#generation-form");
const systemPrompt = element<HTMLTextAreaElement>("#system-prompt");
const promptInput = element<HTMLTextAreaElement>("#prompt");
const precision = element<HTMLSelectElement>("#precision");
const maxNewTokens = element<HTMLInputElement>("#max-new-tokens");
const temperature = element<HTMLInputElement>("#temperature");
const topK = element<HTMLInputElement>("#top-k");
const repetitionPenalty = element<HTMLInputElement>("#repetition-penalty");
const generateButton = element<HTMLButtonElement>("#generate");
const stopButton = element<HTMLButtonElement>("#stop");
const clearButton = element<HTMLButtonElement>("#clear");
const output = element<HTMLPreElement>("#output");
const metric = element<HTMLDivElement>("#metric");
const progressBar = element<HTMLDivElement>("#progress-bar");
const tokenRail = element<HTMLDivElement>("#token-rail");
const diagnosticsOutput = element<HTMLPreElement>("#diagnostics");
const copyDiagnosticsButton = element<HTMLButtonElement>("#copy-diagnostics");

systemPrompt.value = DEFAULT_SYSTEM_PROMPT;
promptInput.value = DEFAULT_PROMPT;
precision.value = DEFAULT_SETTINGS.precision;
maxNewTokens.value = String(DEFAULT_SETTINGS.maxNewTokens);
temperature.value = String(DEFAULT_SETTINGS.temperature);
topK.value = String(DEFAULT_SETTINGS.topK);
repetitionPenalty.value = String(DEFAULT_SETTINGS.repetitionPenalty);

updateWebGpuStatus();
appendDiagnostic("App initialized", clientEnvironment());
setBusy(false);

form.addEventListener("submit", (event) => {
  event.preventDefault();
  startGeneration();
});

stopButton.addEventListener("click", () => {
  jobId += 1;
  const request: WorkerRequest = { type: "reset", jobId };
  worker.postMessage(request);
  setBusy(false);
  metric.textContent = "Stopped";
});

clearButton.addEventListener("click", () => {
  generatedText = "";
  output.textContent = "";
  tokenRail.innerHTML = "";
  diagnostics.length = 0;
  diagnosticsOutput.textContent = "";
  metric.textContent = "Idle";
  progressBar.style.width = "0%";
  appendDiagnostic("Cleared output", clientEnvironment());
});

copyDiagnosticsButton.addEventListener("click", async () => {
  const textToCopy = diagnostics.join("\n");
  await navigator.clipboard.writeText(textToCopy);
  copyDiagnosticsButton.textContent = "Copied";
  window.setTimeout(() => {
    copyDiagnosticsButton.textContent = "Copy log";
  }, 1400);
});

for (const button of document.querySelectorAll<HTMLButtonElement>(".preset")) {
  button.addEventListener("click", () => {
    const preset = PRESETS.find((item) => item.id === button.dataset.preset);
    if (preset) promptInput.value = preset.prompt;
  });
}

worker.addEventListener("message", (event: MessageEvent<WorkerResponse>) => {
  const message = event.data;
  if (message.jobId !== jobId) return;

  switch (message.type) {
    case "status":
      metric.textContent = message.message;
      break;
    case "progress":
      progressBar.style.width = `${message.progress}%`;
      metric.textContent =
        message.progress > 0
          ? `${message.status}: ${message.progress.toFixed(0)}%`
          : `${message.status}: ${message.file}`;
      appendDiagnostic("Worker progress", message);
      break;
    case "debug":
      appendDiagnostic(message.message, message.data);
      break;
    case "token":
      generatedText += message.text;
      output.textContent = generatedText;
      appendTokenMark(message.text);
      output.scrollTop = output.scrollHeight;
      break;
    case "complete":
      generatedText = message.text || generatedText;
      output.textContent = generatedText;
      metric.textContent = `${Math.max(0.1, message.elapsedMs / 1000).toFixed(1)}s`;
      progressBar.style.width = "100%";
      appendDiagnostic("Worker complete", { elapsedMs: message.elapsedMs, textChars: message.text.length });
      setBusy(false);
      break;
    case "error":
      metric.textContent = "Error";
      output.textContent = message.message;
      appendDiagnostic("Worker error", message.data ?? message.message);
      runtimeStatus.classList.add("runtime-error");
      setBusy(false);
      break;
  }
});

function element<T extends Element>(selector: string) {
  const found = document.querySelector<T>(selector);
  if (!found) throw new Error(`Missing element: ${selector}`);
  return found;
}

function readSettings(): GenerationSettings {
  return {
    precision: precision.value as Precision,
    maxNewTokens: readNumber(maxNewTokens, DEFAULT_SETTINGS.maxNewTokens),
    temperature: readNumber(temperature, DEFAULT_SETTINGS.temperature),
    topK: readNumber(topK, DEFAULT_SETTINGS.topK),
    repetitionPenalty: readNumber(repetitionPenalty, DEFAULT_SETTINGS.repetitionPenalty),
  };
}

function readNumber(input: HTMLInputElement, fallback: number) {
  const value = Number(input.value);
  return Number.isFinite(value) ? value : fallback;
}

function startGeneration() {
  if (isGenerating) return;

  generatedText = "";
  output.textContent = "";
  tokenRail.innerHTML = "";
  progressBar.style.width = "0%";
  metric.textContent = "Queueing";
  setBusy(true);
  jobId += 1;

  const request: WorkerRequest = {
    type: "generate",
    jobId,
    systemPrompt: systemPrompt.value.trim() || DEFAULT_SYSTEM_PROMPT,
    prompt: promptInput.value.trim(),
    settings: readSettings(),
  };

  appendDiagnostic("Main generate request", {
    jobId,
    settings: request.settings,
    systemPromptChars: request.systemPrompt.length,
    promptChars: request.prompt.length,
    environment: clientEnvironment(),
  });
  worker.postMessage(request);
}

function setBusy(nextBusy: boolean) {
  isGenerating = nextBusy;
  generateButton.disabled = nextBusy;
  stopButton.disabled = !nextBusy;
  generateButton.textContent = nextBusy ? "Generating" : "Generate";
}

function appendTokenMark(text: string) {
  const mark = document.createElement("span");
  mark.className = "token-mark";
  mark.style.setProperty("--token-size", `${Math.min(42, Math.max(8, text.length * 3))}px`);
  tokenRail.append(mark);
  while (tokenRail.childElementCount > 56) {
    tokenRail.firstElementChild?.remove();
  }
}

function updateWebGpuStatus() {
  const hasWebGpu = "gpu" in navigator;
  runtimeStatus.classList.toggle("runtime-ready", hasWebGpu);
  runtimeStatus.classList.toggle("runtime-error", !hasWebGpu);
  runtimeLabel.textContent = hasWebGpu ? `WebGPU ready / ${MODEL_ID}` : "WebGPU unavailable";
}


worker.addEventListener("error", (event) => {
  appendDiagnostic("Worker script error event", {
    message: event.message,
    filename: event.filename,
    lineno: event.lineno,
    colno: event.colno,
  });
});

window.addEventListener("error", (event) => {
  appendDiagnostic("Window error", {
    message: event.message,
    filename: event.filename,
    lineno: event.lineno,
    colno: event.colno,
    error: event.error instanceof Error ? event.error.stack : String(event.error),
  });
});

window.addEventListener("unhandledrejection", (event) => {
  appendDiagnostic("Window unhandled rejection", {
    reason: event.reason instanceof Error ? event.reason.stack : String(event.reason),
  });
});

function clientEnvironment() {
  return {
    href: window.location.href,
    userAgent: navigator.userAgent,
    hasWebGpu: "gpu" in navigator,
    crossOriginIsolated: window.crossOriginIsolated,
    hardwareConcurrency: navigator.hardwareConcurrency,
    deviceMemory: "deviceMemory" in navigator ? (navigator as Navigator & { deviceMemory?: number }).deviceMemory : undefined,
  };
}

function appendDiagnostic(message: string, data?: unknown) {
  const stamp = new Date().toISOString();
  const suffix = data === undefined ? "" : `\n${safeJson(data)}`;
  diagnostics.push(`[${stamp}] ${message}${suffix}`);
  while (diagnostics.length > 220) diagnostics.shift();
  diagnosticsOutput.textContent = diagnostics.join("\n\n");
  diagnosticsOutput.scrollTop = diagnosticsOutput.scrollHeight;
}

function safeJson(value: unknown) {
  try {
    return JSON.stringify(value, null, 2);
  } catch (error) {
    return String(error);
  }
}