const GPU_PROBE_VERSION = 13;

class ProbeStageError extends Error {
  constructor(message, hung = false, detail) {
    super(message);
    this.hung = hung;
    this.detail = detail;
    this.name = "ProbeStageError";
  }
  hung;
  detail;
}

const DEFAULT_GPU_PROBE_STAGE_TIMEOUT_MS = {
  "navigator-gpu": 1e3,
  "request-adapter": 5e3,
  "adapter-info": 2e3,
  "request-device": 8e3,
  "create-shader-module": 5e3,
  "create-pipeline": 2e4,
  "dispatch-trivial": 1e4,
  "dispatch-arithmetic": 15e3,
  "allocate-storage-buffer": 1e4
};

const PROBE_TRIVIAL_WGSL = `
@group(0) @binding(0) var<storage, read_write> out: array<u32>;

@compute @workgroup_size(1)
fn main() {
  out[0] = 42u;
}
`;
const PROBE_ARITHMETIC_WGSL = `
@group(0) @binding(0) var<storage, read_write> out: array<u32>;

const ITERATIONS: u32 = 4096u;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
  var lo: u32 = id.x * 2654435761u + 1u;
  var hi: u32 = 0u;
  for (var i: u32 = 0u; i < ITERATIONS; i = i + 1u) {
    let a_lo = lo & 0xffffu;
    let a_hi = lo >> 16u;
    let b_lo = 0x9e3779b9u & 0xffffu;
    let b_hi = 0x9e3779b9u >> 16u;
    let mid = a_lo * b_hi + a_hi * b_lo;
    let product_lo = a_lo * b_lo + (mid << 16u);
    let carry = select(0u, 1u, product_lo < a_lo * b_lo);
    hi = a_hi * b_hi + (mid >> 16u) + carry + hi;
    lo = product_lo ^ hi;
  }
  if (id.x == 0u) {
    out[0] = lo ^ hi;
  }
}
`;
const PROBE_ARITHMETIC_WORKGROUPS = 64;

const PROBE_MAP_READ = 1;
const PROBE_COPY_SRC = 4;
const PROBE_COPY_DST = 8;
const PROBE_STORAGE = 128;

const REPORTED_LIMITS = [
  "maxBufferSize",
  "maxStorageBufferBindingSize",
  "maxComputeWorkgroupStorageSize",
  "maxComputeInvocationsPerWorkgroup",
  "maxComputeWorkgroupsPerDimension"
];
const PROBE_ALLOCATION_BYTES = 128 * 1024 * 1024;
const SOFTWARE_ADAPTER_MARKERS = ["swiftshader", "llvmpipe", "lavapipe", "software", "basic render"];
const defaultScheduleTimeout = (callback, ms) => {
  const handle = setTimeout(callback, ms);
  return () => clearTimeout(handle);
};
const describe = (error) => error instanceof Error ? `${error.name}: ${error.message}` : String(error);
function numericLimits(limits) {
  if (!limits) return void 0;
  const collected = {};
  for (const key of REPORTED_LIMITS) {
    const value = limits[key];
    if (typeof value === "number" && Number.isFinite(value)) collected[key] = value;
  }
  return Object.keys(collected).length > 0 ? collected : void 0;
}
function isSoftwareProbeAdapter(info) {
  if (info.isFallbackAdapter === true) return true;
  const identity = [info.vendor, info.architecture, info.device, info.description].filter((part) => typeof part === "string").join(" ").toLowerCase();
  return SOFTWARE_ADAPTER_MARKERS.some((marker) => identity.includes(marker));
}
async function runGpuProbe(options = {}) {
  const now = options.now ?? (() => performance.now());
  const scheduleTimeout = options.scheduleTimeout ?? defaultScheduleTimeout;
  const startedAt = now();
  const breadcrumbs = [];
  const uncapturedErrors = [];
  let deviceLost;
  let deviceReleased = false;
  let adapterReport;
  const record = (breadcrumb) => {
    breadcrumbs.push(breadcrumb);
    options.emit?.(breadcrumb);
  };
  const timeoutFor = (stage2) => options.stageTimeoutMs?.[stage2] ?? DEFAULT_GPU_PROBE_STAGE_TIMEOUT_MS[stage2];
  const stage = async (name, run, detailOf) => {
    record({ stage: name, status: "started", atMs: now() - startedAt });
    const stageStartedAt = now();
    let cancelTimeout;
    try {
      const limit = timeoutFor(name);
      const value = await new Promise((resolve, reject) => {
        cancelTimeout = scheduleTimeout(() => {
          reject(new ProbeStageError(`${name} did not return within ${limit} ms`, true));
        }, limit);
        Promise.resolve().then(run).then(resolve, (error) => {
          reject(error instanceof ProbeStageError ? error : new ProbeStageError(describe(error)));
        });
      });
      record({
        stage: name,
        status: "passed",
        atMs: now() - startedAt,
        durationMs: now() - stageStartedAt,
        detail: detailOf?.(value)
      });
      return value;
    } catch (error) {
      const failure = error instanceof ProbeStageError ? error : new ProbeStageError(describe(error));
      record({
        stage: name,
        status: "failed",
        atMs: now() - startedAt,
        durationMs: now() - stageStartedAt,
        message: failure.message,
        detail: failure.detail
      });
      throw failure;
    } finally {
      cancelTimeout?.();
    }
  };
  const finish = (outcome, failedStage, reason) => ({
    outcome,
    failedStage,
    reason,
    adapter: adapterReport,
    deviceLost,
    uncapturedErrors: uncapturedErrors.length > 0 ? uncapturedErrors : void 0,
    breadcrumbs,
    totalMs: now() - startedAt,
    completedAt: options.clock?.(),
    probeVersion: GPU_PROBE_VERSION
  });
  let device;
  try {
    const gpu = await stage("navigator-gpu", () => {
      const candidate = options.gpu ?? globalThis.navigator?.gpu;
      if (!candidate || typeof candidate.requestAdapter !== "function") {
        throw new ProbeStageError("navigator.gpu is not available in this context");
      }
      return candidate;
    });
    const adapter = await stage("request-adapter", async () => {
      const requested = await gpu.requestAdapter({ powerPreference: "high-performance" });
      if (!requested) throw new ProbeStageError("requestAdapter returned no adapter");
      return requested;
    });
    adapterReport = await stage(
      "adapter-info",
      async () => {
        const info = adapter.info ?? await adapter.requestAdapterInfo?.() ?? {};
        const collected = {
          vendor: info.vendor,
          architecture: info.architecture,
          device: info.device,
          description: info.description,
          isFallbackAdapter: adapter.isFallbackAdapter,
          limits: numericLimits(adapter.limits),
          features: adapter.features ? [...adapter.features] : void 0
        };
        if (isSoftwareProbeAdapter(collected)) {
          throw new ProbeStageError(
            "WebGPU resolved to a software adapter, which is slower than the CPU prover",
            false,
            { ...collected }
          );
        }
        return collected;
      },
      (value) => ({ ...value })
    );
    device = await stage("request-device", async () => {
      const created = await adapter.requestDevice();
      created.addEventListener?.("uncapturederror", (event) => {
        uncapturedErrors.push(event.error?.message ?? "unknown uncaptured GPU error");
      });
      void created.lost?.then((lost) => {
        if (lost && !deviceReleased) {
          deviceLost = `${lost.reason ?? "unknown"}: ${lost.message ?? ""}`.trim();
        }
      }, () => void 0);
      return created;
    });
    const gpuDevice = device;
    const shaderModule = await stage(
      "create-shader-module",
      async () => {
        const created = gpuDevice.createShaderModule({ code: PROBE_TRIVIAL_WGSL });
        const info = await created.getCompilationInfo?.();
        const errors = (info?.messages ?? []).filter((message) => message.type === "error");
        if (errors.length > 0) {
          throw new ProbeStageError(
            `WGSL compilation reported ${errors.length} error(s): ${errors[0]?.message ?? ""}`,
            false,
            { messages: errors.map((message) => message.message) }
          );
        }
        return created;
      }
    );
    const pipeline = await stage(
      "create-pipeline",
      () => gpuDevice.createComputePipelineAsync({
        layout: "auto",
        compute: { module: shaderModule, entryPoint: "main" }
      })
    );
    const dispatch = async (target, workgroups) => {
      const outputBytes = 4;
      const output = gpuDevice.createBuffer({
        size: outputBytes,
        usage: PROBE_STORAGE | PROBE_COPY_SRC
      });
      const readback = gpuDevice.createBuffer({
        size: outputBytes,
        usage: PROBE_COPY_DST | PROBE_MAP_READ
      });
      try {
        const bindGroup = gpuDevice.createBindGroup({
          layout: target.getBindGroupLayout(0),
          entries: [{ binding: 0, resource: { buffer: output } }]
        });
        const encoder = gpuDevice.createCommandEncoder();
        const pass = encoder.beginComputePass();
        pass.setPipeline(target);
        pass.setBindGroup(0, bindGroup);
        pass.dispatchWorkgroups(workgroups);
        pass.end();
        encoder.copyBufferToBuffer(output, 0, readback, 0, outputBytes);
        gpuDevice.queue.submit([encoder.finish()]);
        await readback.mapAsync(PROBE_MAP_READ);
        const value = new Uint32Array(readback.getMappedRange().slice(0))[0] ?? 0;
        readback.unmap();
        return value;
      } finally {
        output.destroy?.();
        readback.destroy?.();
      }
    };
    await stage(
      "dispatch-trivial",
      async () => {
        const value = await dispatch(pipeline, 1);
        if (value !== 42) {
          throw new ProbeStageError(
            `trivial kernel wrote ${value}, expected 42`,
            false,
            { value }
          );
        }
        return value;
      }
    );
    await stage("dispatch-arithmetic", async () => {
      const module = gpuDevice.createShaderModule({ code: PROBE_ARITHMETIC_WGSL });
      const arithmetic = await gpuDevice.createComputePipelineAsync({
        layout: "auto",
        compute: { module, entryPoint: "main" }
      });
      await dispatch(arithmetic, PROBE_ARITHMETIC_WORKGROUPS);
      return void 0;
    });
    await stage(
      "allocate-storage-buffer",
      async () => {
        const limits = numericLimits(gpuDevice.limits) ?? adapterReport?.limits ?? {};
        const ceiling = Math.min(
          limits.maxBufferSize ?? PROBE_ALLOCATION_BYTES,
          limits.maxStorageBufferBindingSize ?? PROBE_ALLOCATION_BYTES,
          PROBE_ALLOCATION_BYTES
        );
        const size = Math.max(4, Math.floor(ceiling / 4) * 4);
        let buffer;
        try {
          buffer = gpuDevice.createBuffer({ size, usage: PROBE_STORAGE | PROBE_COPY_DST });
          gpuDevice.queue.writeBuffer(buffer, 0, new Uint32Array([1, 2, 3, 4]));
          gpuDevice.queue.submit([]);
          await gpuDevice.queue.onSubmittedWorkDone();
        } finally {
          buffer?.destroy?.();
        }
        return { size };
      },
      (value) => ({ ...value })
    );
    if (deviceLost) {
      return finish(
        "failed",
        breadcrumbs.at(-1)?.stage,
        `WebGPU device was lost during the probe (${deviceLost})`
      );
    }
    return finish("ok");
  } catch (error) {
    const failure = error instanceof ProbeStageError ? error : void 0;
    const failedStage = breadcrumbs.at(-1)?.stage;
    const message = failure?.message ?? describe(error);
    if (failedStage === "navigator-gpu") return finish("unsupported", failedStage, message);
    return finish(failure?.hung === true ? "hung" : "failed", failedStage, message);
  } finally {
    deviceReleased = true;
    device?.destroy?.();
  }
}

const post = (message) => {
  self.postMessage(message);
};
self.addEventListener("message", (event) => {
  if (event.data?.type !== "run") return;
  void runGpuProbe({
    stageTimeoutMs: event.data.stageTimeoutMs,
    clock: () => Date.now(),
    emit: (breadcrumb) => {
      post({ type: "breadcrumb", breadcrumb });
    }
  }).then((report) => {
    post({ type: "report", report });
  });
});
