/**
 * The sandbox page: editor on the left, output on the right, a timing table
 * underneath that lines up GPU tic/toc segments with the CPU run's.
 */
import { CodeEditor } from './editor/codeEditor.ts';
import { requestSandboxDevice, runScript, type SandboxGpu } from './mgpu/session.ts';
import { formatFailure } from './mgpu/errors.ts';
import type { CpuRunHandle } from './cpu/cpuRunner.ts';
import fusedLoop from '../examples/fused_loop.m?raw';
import matmul from '../examples/matmul.m?raw';
import monteCarloPi from '../examples/monte_carlo_pi.m?raw';
import logisticEnsemble from '../examples/logistic_ensemble.m?raw';
import reductions from '../examples/reductions.m?raw';

const EXAMPLES: { key: string; title: string; source: string }[] = [
  { key: 'fused_loop', title: 'fused elementwise loop', source: fusedLoop },
  { key: 'matmul', title: 'matrix multiply (GEMM)', source: matmul },
  { key: 'monte_carlo_pi', title: 'Monte Carlo pi (masks)', source: monteCarloPi },
  { key: 'logistic_ensemble', title: 'logistic map ensemble', source: logisticEnsemble },
  { key: 'reductions', title: 'fused reductions', source: reductions },
];

const $ = <T extends HTMLElement>(id: string): T => {
  const el = document.getElementById(id);
  if (!el) throw new Error(`missing #${id}`);
  return el as T;
};

const consoleEl = $<HTMLPreElement>('console');
const cpuConsoleEl = $<HTMLPreElement>('cpuconsole');
const runBtn = $<HTMLButtonElement>('run');
const runCpuBtn = $<HTMLButtonElement>('runcpu');
const copyBtn = $<HTMLButtonElement>('copy');
const exampleSel = $<HTMLSelectElement>('example');
const timingsEl = $<HTMLTableElement>('timings');
const planDetails = $<HTMLElement>('plandetails');
const planEl = $<HTMLPreElement>('plan');
const runState = $<HTMLElement>('runstate');
const deviceEl = $<HTMLElement>('device');
const gpuWarn = $<HTMLElement>('gpuwarn');

const editor = new CodeEditor({
  textarea: $<HTMLTextAreaElement>('source'),
  overlay: $('highlight'),
  onInput: () => {
    $('editstate').textContent = '';
  },
});

for (const ex of EXAMPLES) {
  const opt = document.createElement('option');
  opt.value = ex.key;
  opt.textContent = ex.title;
  exampleSel.appendChild(opt);
}
exampleSel.addEventListener('change', () => {
  const ex = EXAMPLES.find((e) => e.key === exampleSel.value);
  if (ex) editor.value = ex.source;
});
editor.value = EXAMPLES[0].source;

const append = (el: HTMLElement, text: string, cls?: string): void => {
  if (cls) {
    const span = document.createElement('span');
    span.className = cls;
    span.textContent = text;
    el.appendChild(span);
  } else {
    el.appendChild(document.createTextNode(text));
  }
  el.scrollTop = el.scrollHeight;
};

/** Latest timing results, merged into the table by tic..toc pair. */
const lastSeconds: { gpu: number[]; cpu: number[] } = { gpu: [], cpu: [] };
let gpuTotal: number | null = null;
let cpuTotal: number | null = null;

function renderTimings(): void {
  const n = Math.max(lastSeconds.gpu.length, lastSeconds.cpu.length);
  if (n === 0 && gpuTotal === null && cpuTotal === null) {
    timingsEl.style.display = 'none';
    return;
  }
  const fmt = (s: number | undefined | null): string =>
    s === undefined || s === null ? '—' : s >= 0.1 ? `${s.toFixed(3)} s` : `${(s * 1000).toFixed(2)} ms`;
  const rows: string[] = [
    '<tr><th></th><th>GPU</th><th>CPU (numbl)</th><th>CPU / GPU</th></tr>',
  ];
  for (let i = 0; i < n; i++) {
    const g = lastSeconds.gpu[i];
    const c = lastSeconds.cpu[i];
    const ratio = g !== undefined && c !== undefined ? `${(c / g).toFixed(1)}×` : '—';
    rows.push(
      `<tr><td>tic…toc #${i + 1}</td><td>${fmt(g)}</td><td>${fmt(c)}</td><td>${ratio}</td></tr>`,
    );
  }
  const totalRatio =
    gpuTotal !== null && cpuTotal !== null ? `${(cpuTotal / gpuTotal).toFixed(1)}×` : '—';
  rows.push(
    `<tr><td>whole script</td><td>${fmt(gpuTotal)}</td><td>${fmt(cpuTotal)}</td><td>${totalRatio}</td></tr>`,
  );
  timingsEl.innerHTML = rows.join('');
  timingsEl.style.display = 'table';
}

let gpu: SandboxGpu | null = null;
(async () => {
  try {
    gpu = await requestSandboxDevice();
    deviceEl.textContent = `GPU: ${gpu.description}`;
    gpu.device.lost.then((info) => {
      gpu = null;
      gpuWarn.hidden = false;
      gpuWarn.textContent = `The GPU device was lost (${info.message}); reload the page.`;
      runBtn.disabled = true;
    });
  } catch (e) {
    gpuWarn.hidden = false;
    gpuWarn.innerHTML =
      `<b>No WebGPU here.</b> ${e instanceof Error ? e.message : String(e)} — ` +
      `Chrome/Edge have it on by default; Firefox and Safari are rolling it out. ` +
      `The CPU run still works.`;
    runBtn.disabled = true;
  }
})();

runBtn.addEventListener('click', () => {
  void runGpu();
});
async function runGpu(): Promise<void> {
  if (!gpu) return;
  runBtn.disabled = true;
  runState.textContent = 'compiling…';
  consoleEl.textContent = '';
  lastSeconds.gpu = [];
  gpuTotal = null;
  renderTimings();
  planDetails.hidden = true;
  const source = editor.value;
  try {
    const t0 = performance.now();
    const run = await runScript(gpu.device, source, (text) => {
      runState.textContent = 'running…';
      append(consoleEl, text);
    });
    void t0;
    if (run.result.error) {
      append(consoleEl, `\n${run.result.error}\n`, 'err');
    }
    append(
      consoleEl,
      `\n[compile ${(run.compileSeconds * 1000).toFixed(0)} ms · ` +
        `run ${run.result.totalSeconds.toFixed(3)} s · f32 · ${gpu.description}]\n`,
      'meta',
    );
    lastSeconds.gpu = run.result.segments.map((s) => s.seconds);
    gpuTotal = run.result.totalSeconds;
    planEl.textContent = run.planDescription.join('\n') || '(no GPU ops)';
    planDetails.hidden = false;
    renderTimings();
    runState.textContent = '';
  } catch (e) {
    append(consoleEl, formatFailure(e, source) + '\n', 'err');
    runState.textContent = '';
  } finally {
    runBtn.disabled = !gpu;
  }
}

let cpuHandle: CpuRunHandle | null = null;
runCpuBtn.addEventListener('click', () => {
  if (cpuHandle) {
    cpuHandle.cancel();
    cpuHandle = null;
    runCpuBtn.textContent = 'Run on CPU (numbl)';
    return;
  }
  void runCpu();
});
async function runCpu(): Promise<void> {
  cpuConsoleEl.style.display = 'block';
  cpuConsoleEl.textContent = '';
  append(cpuConsoleEl, '[CPU · numbl engine in a worker · f64]\n', 'meta');
  lastSeconds.cpu = [];
  cpuTotal = null;
  renderTimings();
  runCpuBtn.textContent = 'Stop CPU run';
  // numbl's browser engine is a couple of megabytes; load it on first use.
  const { runOnCpu } = await import('./cpu/cpuRunner.ts');
  const handle = runOnCpu(editor.value, (text) => append(cpuConsoleEl, text));
  cpuHandle = handle;
  try {
    const res = await handle.result;
    if (res.error) append(cpuConsoleEl, `\n${res.error}\n`, 'err');
    if (res.aborted) append(cpuConsoleEl, `\n[stopped]\n`, 'meta');
    else {
      append(cpuConsoleEl, `[total ${res.totalSeconds.toFixed(3)} s]\n`, 'meta');
      // Line the CPU's tic..toc pairs up with the GPU's by order.
      lastSeconds.cpu = [...res.output.matchAll(/Elapsed time is ([0-9.eE+-]+) seconds/g)]
        .map((m) => Number(m[1]));
      cpuTotal = res.totalSeconds;
      renderTimings();
    }
  } catch (e) {
    append(cpuConsoleEl, `${e instanceof Error ? e.message : String(e)}\n`, 'err');
  } finally {
    if (cpuHandle === handle) cpuHandle = null;
    runCpuBtn.textContent = 'Run on CPU (numbl)';
  }
}

copyBtn.addEventListener('click', () => {
  void navigator.clipboard.writeText(editor.value).then(() => {
    copyBtn.textContent = 'Copied ✓';
    setTimeout(() => {
      copyBtn.textContent = 'Copy script';
    }, 1200);
  });
});