import { alignTo, toTypedArray } from "./utils.js";
import { debugLog, trackResource, untrackResource } from "./debug.js";
import { nextId } from "./shared.js";
import type {
  BufferData,
  BufferUsagePreset,
  GpuBuffer,
  GpuContext,
  TypedArray,
} from "./types.js";

/**
 * Type guard for {@link GpuBuffer} resource objects (raw GPUBuffers also have a
 * similar shape — discriminate on `id`).
 */
export const isGpuBuffer = (value: unknown): value is GpuBuffer =>
  typeof value === "object" &&
  value !== null &&
  "id" in value &&
  "buffer" in value;

/**
 * Usage presets for {@link createBuffer}: common GPUBufferUsage flag
 * combinations by name.
 */
export const BUFFER_USAGE_PRESETS: Record<
  BufferUsagePreset,
  GPUBufferUsageFlags
> = {
  // Draw inputs
  vertex: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
  index: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
  // STORAGE so compute shaders can write draw arguments (GPU-driven drawing)
  indirect:
    GPUBufferUsage.INDIRECT | GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  // Bindable buffers, named after GPUBufferBindingType:
  // uniform → var<uniform>, storage → var<storage, read_write> (COPY_SRC for
  // readBuffer), read-only-storage → var<storage, read>
  uniform: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  storage:
    GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
  "read-only-storage": GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  // Staging buffers, named after D3D12 heap types
  upload: GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC,
  readback: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
  // Timestamp/occlusion query resolve target
  "query-resolve": GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
};

/**
 * Coerce input to a typed array, additionally converting index data to one of
 * the two WebGPU index formats: libraries like primitive-geometry hand out
 * Uint8Array cells for small geometries, which uint16/uint32 draws misread. An
 * explicit `indexFormat` is authoritative; otherwise it is derived from the
 * data (uint32 when any index exceeds 65535). Raw ArrayBuffers are trusted to
 * already contain data in the right format.
 */
function toBufferData(
  data: BufferData,
  index: boolean,
  indexFormat?: GPUIndexFormat,
): TypedArray {
  const typed = toTypedArray(data, { index });
  if (!index || data instanceof ArrayBuffer) return typed;

  const values = typed as ArrayLike<number>;
  const format =
    indexFormat ??
    (typed instanceof Uint32Array
      ? "uint32"
      : typed instanceof Uint16Array
        ? "uint16"
        : Array.prototype.some.call(values, (value: number) => value > 0xff_ff)
          ? "uint32"
          : "uint16");

  if (format === "uint32") {
    return typed instanceof Uint32Array ? typed : new Uint32Array(values);
  }
  return typed instanceof Uint16Array ? typed : new Uint16Array(values);
}

export interface CreateBufferOptions extends GPUObjectDescriptorBase {
  /** Usage preset name, or raw GPUBufferUsage flags. */
  usage: BufferUsagePreset | GPUBufferUsageFlags;
  /**
   * Initial contents. Plain arrays coerce to Float32Array (Uint16/Uint32Array
   * for index data).
   */
  data?: BufferData;
  /** Byte size when allocating without data, or to over-allocate. */
  size?: number;
  /**
   * Index format for index buffers. Derived from the data when omitted;
   * required knowledge for buffers allocated by size and filled later, as
   * updateBuffer converts index data to the buffer's declared format.
   */
  indexFormat?: GPUIndexFormat;
}

/**
 * Create a GPU buffer.
 *
 * ```js
 * const positions = gpu.createBuffer(ctx, {
 *   usage: "vertex",
 *   data: geometry.positions,
 * });
 * const indices = gpu.createBuffer(ctx, {
 *   usage: "index",
 *   data: geometry.cells,
 * });
 * ```
 */
export function createBuffer(
  ctx: GpuContext,
  options: CreateBufferOptions,
): GpuBuffer {
  const usage =
    typeof options.usage === "string"
      ? BUFFER_USAGE_PRESETS[options.usage]
      : options.usage;
  const index = (usage & GPUBufferUsage.INDEX) !== 0;
  const data =
    options.data === undefined
      ? undefined
      : toBufferData(options.data, index, options.indexFormat);

  const size = alignTo(Math.max(options.size ?? 0, data?.byteLength ?? 0), 4);
  if (size === 0) {
    throw new Error("pex-gpu: createBuffer needs data or a size");
  }

  // mappedAtCreation sidesteps writeBuffer's multiple-of-4 size requirement
  // (the buffer size is 4-aligned; the data doesn't have to be, e.g. an odd
  // number of uint16 indices)
  const buffer = ctx.device.createBuffer({
    size,
    usage,
    ...(data && { mappedAtCreation: true }),
    ...(options.label && { label: options.label }),
  });
  if (data) {
    new Uint8Array(buffer.getMappedRange()).set(
      new Uint8Array(data.buffer, data.byteOffset, data.byteLength),
    );
    buffer.unmap();
  }

  const id = nextId();
  trackResource(ctx.device, "buffers", {
    id,
    ...(options.label && { label: options.label }),
    usage: buffer.usage,
    bytes: size,
  });
  debugLog(
    ctx.device,
    "resources",
    () =>
      `createBuffer ${options.label ? `"${options.label}" ` : ""}${size} bytes`,
  );
  let disposed = false;
  function dispose(): void {
    if (disposed) return;
    disposed = true;
    untrackResource(ctx.device, "buffers", id);
    buffer.destroy();
  }

  return {
    buffer,
    size,
    usage: buffer.usage,
    id,
    ...(data && { length: data.length }),
    ...(index && {
      indexFormat:
        options.indexFormat ??
        ((data instanceof Uint32Array ? "uint32" : "uint16") as GPUIndexFormat),
    }),
    dispose,
    [Symbol.dispose]: dispose,
  };
}

/**
 * Upload data into an existing buffer.
 *
 * ```js
 * gpu.updateBuffer(ctx, positions, newPositions);
 * ```
 */
export function updateBuffer(
  ctx: GpuContext,
  target: GpuBuffer,
  data: BufferData,
  byteOffset = 0,
): void {
  // Index data converts to the buffer's declared format so they cannot drift
  const typed = toBufferData(
    data,
    (target.usage & GPUBufferUsage.INDEX) !== 0,
    target.indexFormat,
  );
  // Views here are never SharedArrayBuffer-backed (writeBuffer rejects those)
  let bytes = new Uint8Array(
    typed.buffer as ArrayBuffer,
    typed.byteOffset,
    typed.byteLength,
  );
  // writeBuffer sizes must be multiples of 4 — zero-pad the tail
  if (bytes.byteLength % 4 !== 0) {
    const padded = new Uint8Array(alignTo(bytes.byteLength, 4));
    padded.set(bytes);
    bytes = padded;
  }
  ctx.device.queue.writeBuffer(target.buffer, byteOffset, bytes);

  // `length` describes one array's worth of contents, so only a write that
  // replaces them restates it: a buffer allocated by size and filled in chunks
  // never had it set, and must not gain the length of whichever chunk landed
  // at offset 0.
  if (byteOffset === 0 && target.length !== undefined) {
    target.length = typed.length;
  }
}