import type { BackendInfo, ComputeInfo, RunRequest, RunResult } from "./types";
import { executeBundle, formatError } from "./execute";
import { LIBRARY_FILES } from "../generate/library";
import { runFdLoopGpu, type GpuHandle } from "../gpu/fdLoop";
import { fdLoopRecordToPlain, fdLoopSetupFromPlain } from "../gpu/fdStage";

/** Supplies the (cached) WebGPU device of the host, or rejects when there is none. */
export type GpuProvider = () => Promise<GpuHandle>;

/**
 * Run a request, on the GPU when it asks for that and the host has one. The
 * plain path is executeBundle. The staged path runs the setup script in numbl,
 * the loop on the GPU and the output script in numbl, and reports where the
 * loop ran in `compute`; when the GPU was asked for but cannot be used, the
 * whole bundle runs in numbl and `compute.note` says why.
 */
export async function executeRequest(
  req: RunRequest,
  backend: BackendInfo,
  onOutput?: (text: string) => void,
  gpu?: GpuProvider,
): Promise<RunResult> {
  const stages = req.gpu;
  if (!stages) return executeBundle(req, backend, onOutput);
  const fallback = (note: string): RunResult => ({
    ...executeBundle(req, backend, onOutput),
    compute: { loop: "numbl", note },
  });

  if (!gpu) return fallback("WebGPU is not available here, so the time loop ran in numbl");
  const shipped = LIBRARY_FILES.find((f) => f.name === stages.loopFile)?.source;
  const current = req.bundle.files.find((f) => f.name === stages.loopFile)?.source;
  if (shipped !== undefined && current !== shipped)
    return fallback(
      `${stages.loopFile} was edited, so the time loop ran in numbl, where the edit is live`,
    );
  let handle: GpuHandle;
  try {
    handle = await gpu();
  } catch (e) {
    return fallback(`no WebGPU device (${formatError(e)}), so the time loop ran in numbl`);
  }

  const t0 = performance.now();
  const output: string[] = [];
  const say = (text: string) => {
    output.push(text);
    onOutput?.(text);
  };
  const libs = req.bundle.files.filter((f) => f.name !== req.bundle.entry);
  const fail = (error: string, compute?: ComputeInfo): RunResult => ({
    ok: false,
    error,
    output,
    wallMs: performance.now() - t0,
    backend,
    compute,
  });

  const setupRun = executeBundle(
    {
      bundle: {
        files: [{ name: "gpu_setup.m", source: stages.setup }, ...libs],
        entry: "gpu_setup.m",
        resultVariable: "S",
      },
    },
    backend,
    say,
  );
  if (!setupRun.ok) return fail(setupRun.error);
  let setup;
  try {
    setup = fdLoopSetupFromPlain(setupRun.value);
  } catch (e) {
    return fail(`the setup stage returned an unexpected shape: ${formatError(e)}`);
  }

  let lastPct = 0;
  let rec;
  try {
    rec = await runFdLoopGpu(handle.device, setup, (done, total) => {
      const pct = Math.floor((10 * done) / total);
      if (pct > lastPct && done < total) {
        lastPct = pct;
        say(`gpu: ${done.toLocaleString()} of ${total.toLocaleString()} steps\n`);
      }
    });
  } catch (e) {
    return fail(`WebGPU loop: ${formatError(e)}`, { loop: "webgpu", adapter: handle.adapter });
  }

  const Rec = fdLoopRecordToPlain(rec, setup, handle.adapter);
  const outRun = executeBundle(
    {
      bundle: {
        files: [{ name: "gpu_output.m", source: stages.output }, ...libs],
        entry: "gpu_output.m",
        resultVariable: req.bundle.resultVariable,
      },
      inputs: { ...(req.inputs ?? {}), Rec },
    },
    backend,
    say,
  );
  const compute: ComputeInfo = {
    loop: "webgpu",
    adapter: handle.adapter,
    gpuMs: rec.gpuMs,
    numblMs: setupRun.wallMs + outRun.wallMs,
  };
  if (!outRun.ok) return fail(outRun.error, compute);
  return {
    ok: true,
    value: outRun.value,
    output,
    wallMs: performance.now() - t0,
    backend,
    compute,
  };
}