import { BUFFER_USAGE_PRESETS } from "./buffer.js";
import { trackResource, untrackResource } from "./debug.js";
import { nextId } from "./shared.js";
import type { GpuContext } from "./types.js";

export interface CreateTimestampQueryOptions {
  /** Timestamp slots: two per pass timed (begin/end). Defaults to 2. */
  count?: number;
  /**
   * Readback buffers. One (the default) reads a frame only when the previous
   * read has landed, skipping the others; three is what it takes to read every
   * frame with two frames in flight.
   */
  buffering?: number;
  label?: string;
}

export interface TimestampQuery extends Disposable {
  readonly querySet: GPUQuerySet;
  /** Number of timestamp slots (2 = begin/end of pass). */
  readonly count: number;
  /** Pass to a render/compute pass descriptor to time the whole pass. */
  readonly timestampWrites: GPURenderPassTimestampWrites;
  /**
   * Encode resolution of the first `count` queries into a free readback buffer,
   * and arm it for the next {@link TimestampQuery.read}. Call once per frame,
   * after the passes being timed. False when every buffer is still mapping —
   * that frame goes unread rather than stalling on one. Resolving twice without
   * reading in between overwrites the armed buffer.
   */
  resolve(encoder: GPUCommandEncoder, count?: number): boolean;
  /**
   * Read the armed buffer's timestamps, in nanoseconds. Null when nothing is
   * armed or the device went away, so it can be polled every frame. Only call
   * after the encoder passed to `resolve()` has actually been submitted (eg.
   * after `endFrame()`, or fire-and-forget from inside a `frame()` callback) —
   * awaiting it before that submit stalls the very submit it depends on and
   * silently returns stale data instead of the intended timestamps.
   */
  read(): Promise<BigInt64Array | null>;
  /**
   * Release the query set and readback buffers. Identical to
   * `[Symbol.dispose]`, for call sites without `using`. Safe to call with a
   * read in flight: the GPU objects go once it settles.
   */
  dispose(): void;
}

interface ReadbackSlot {
  resolve: GPUBuffer;
  read: GPUBuffer;
  resolveId: number;
  readId: number;
  /** Armed by resolve(), released once read() settles. */
  busy: boolean;
  /** Queries this slot holds, which may be fewer than the set's count. */
  count: number;
}

/**
 * GPU timestamp measurement. Requires creating the context with
 * `requiredFeatures: ["timestamp-query"]`.
 *
 * One query set plus a ring of readback buffers; `debug(ctx, { timing: true })`
 * builds on the same object to time every pass automatically.
 *
 * ```js
 * const query = gpu.createTimestampQuery(ctx);
 * // pass descriptor: { ..., timestampWrites: query.timestampWrites }
 * query.resolve(encoder);
 * const timestamps = await query.read();
 * if (timestamps)
 *   console.log(`${Number(timestamps[1] - timestamps[0]) / 1e6} ms`);
 * ```
 */
export function createTimestampQuery(
  ctx: GpuContext,
  options: CreateTimestampQueryOptions = {},
): TimestampQuery {
  if (!ctx.device.features.has("timestamp-query")) {
    throw new Error(
      'pex-gpu: createTimestampQuery requires the "timestamp-query" feature; pass requiredFeatures: ["timestamp-query"] to createContext',
    );
  }
  const { count = 2, label = "pex-gpu timestamps" } = options;
  const buffering = Math.max(1, options.buffering ?? 1);

  const size = count * 8;
  const querySet = ctx.device.createQuerySet({
    label,
    type: "timestamp",
    count,
  });
  const querySetId = nextId();
  trackResource(ctx.device, "querySets", {
    id: querySetId,
    label: querySet.label,
    type: "timestamp",
    count,
    bytes: size,
  });

  const slots: ReadbackSlot[] = [];
  for (let index = 0; index < buffering; index++) {
    const suffix = buffering > 1 ? ` ${index}` : "";
    const resolve = ctx.device.createBuffer({
      label: `${label} resolve${suffix}`,
      size,
      usage: BUFFER_USAGE_PRESETS["query-resolve"],
    });
    const read = ctx.device.createBuffer({
      label: `${label} read${suffix}`,
      size,
      usage: BUFFER_USAGE_PRESETS.readback,
    });
    const resolveId = nextId();
    const readId = nextId();
    for (const [id, buffer] of [
      [resolveId, resolve],
      [readId, read],
    ] as const) {
      trackResource(ctx.device, "buffers", {
        id,
        label: buffer.label,
        usage: buffer.usage,
        bytes: size,
      });
    }
    slots.push({ resolve, read, resolveId, readId, busy: false, count });
  }

  let armed: ReadbackSlot | null = null;
  let mapping = 0;
  let disposed = false;

  function destroy(): void {
    untrackResource(ctx.device, "querySets", querySetId);
    querySet.destroy();
    for (const slot of slots) {
      untrackResource(ctx.device, "buffers", slot.resolveId);
      untrackResource(ctx.device, "buffers", slot.readId);
      slot.resolve.destroy();
      slot.read.destroy();
    }
  }

  function dispose(): void {
    if (disposed) return;
    disposed = true;
    // A buffer being mapped can't be destroyed; the last read() takes over.
    if (!mapping) destroy();
  }

  return {
    querySet,
    count,
    timestampWrites: {
      querySet,
      beginningOfPassWriteIndex: 0,
      endOfPassWriteIndex: 1,
    },
    resolve(encoder, resolveCount = count) {
      const slot = armed ?? slots.find((slot) => !slot.busy);
      if (!slot) return false;
      slot.busy = true;
      slot.count = resolveCount;
      armed = slot;
      encoder.resolveQuerySet(querySet, 0, resolveCount, slot.resolve, 0);
      encoder.copyBufferToBuffer(
        slot.resolve,
        0,
        slot.read,
        0,
        resolveCount * 8,
      );
      return true;
    },
    async read() {
      const slot = armed;
      armed = null;
      if (!slot) return null;
      mapping++;
      try {
        // Wait a macrotask, not just a microtask: frame()'s tick does
        // `await callback(info)` before endFrame()'s queue.submit(), which is
        // itself only one microtask hop — a plain `await null` here can
        // resolve first (eg. when the callback fires off read() without
        // awaiting it) and call mapAsync() before that submit lands, which
        // WebGPU rejects ("used in submit while mapped"). A macrotask
        // boundary is only reached once the whole microtask queue, including
        // endFrame's submit, has drained.
        await new Promise((resolve) => setTimeout(resolve, 0));
        await slot.read.mapAsync(GPUMapMode.READ);
        // Copy out before unmap() detaches the mapped range.
        const timestamps = new BigInt64Array(
          slot.read.getMappedRange(0, slot.count * 8).slice(0),
        );
        slot.read.unmap();
        return timestamps;
      } catch {
        return null; // device lost, or disposed mid-read
      } finally {
        slot.busy = false;
        mapping--;
        if (disposed && !mapping) destroy();
      }
    },
    dispose,
    [Symbol.dispose]: dispose,
  };
}