import { BUFFER_USAGE_PRESETS } from "./buffer.js";
import { trackResource, untrackResource } from "./debug.js";
import { nextId } from "./shared.js";
import {
  alignTo,
  assertBlockAlignedOrigin,
  physicalExtent,
  texelCopyLayout,
} from "./utils.js";
import type { GpuBuffer, GpuContext, GpuTexture } from "./types.js";

export interface ReadTextureOptions extends Pick<
  GPUTexelCopyTextureInfo,
  "mipLevel"
> {
  /** Region origin. Defaults to 0. */
  x?: number;
  /** Region origin. Defaults to 0. */
  y?: number;
  /** Region size. Defaults to the mip level's width. */
  width?: number;
  /** Region size. Defaults to the mip level's height. */
  height?: number;
  /** Array layer or cube face to read from. */
  layer?: number;
}

/**
 * Asynchronously read pixels back from a texture (eg. for picking). The source
 * needs a `usage` including COPY_SRC. Compressed formats read back as their raw
 * block bytes, the region rounded out to whole blocks.
 *
 * ```js
 * const [r, g, b, a] = await gpu.readTexture(ctx, pickingTexture, {
 *   x,
 *   y,
 *   width: 1,
 *   height: 1,
 * });
 * ```
 */
export async function readTexture(
  ctx: GpuContext,
  source: GpuTexture,
  options: ReadTextureOptions = {},
): Promise<Uint8Array> {
  const mipLevel = options.mipLevel ?? 0;
  const {
    x = 0,
    y = 0,
    width = Math.max(1, source.width >> mipLevel),
    height = Math.max(1, source.height >> mipLevel),
  } = options;
  const origin: GPUOrigin3D = [x, y, options.layer ?? 0];
  assertBlockAlignedOrigin(source.format, origin);

  const { bytesPerRow: rowBytes, rowsPerImage } = texelCopyLayout(
    source.format,
    width,
    height,
  );
  const bytesPerRow = alignTo(rowBytes, 256);

  const readback = ctx.device.createBuffer({
    label: "pex-gpu readback",
    size: bytesPerRow * rowsPerImage,
    usage: BUFFER_USAGE_PRESETS.readback,
  });
  const readbackId = nextId();
  trackResource(ctx.device, "buffers", {
    id: readbackId,
    label: readback.label,
    usage: readback.usage,
    bytes: readback.size,
  });

  const encoder = ctx.device.createCommandEncoder({
    label: "pex-gpu readTexture",
  });
  encoder.copyTextureToBuffer(
    { texture: source.texture, mipLevel, origin },
    { buffer: readback, bytesPerRow, rowsPerImage },
    physicalExtent(source.format, width, height),
  );
  ctx.device.queue.submit([encoder.finish()]);

  await readback.mapAsync(GPUMapMode.READ);
  const mapped = new Uint8Array(readback.getMappedRange());
  const data = new Uint8Array(rowBytes * rowsPerImage);
  // Buffer copies require bytesPerRow to be a 256-byte multiple, so mapped
  // rows carry trailing padding — copy row by row into a tightly packed buffer.
  for (let row = 0; row < rowsPerImage; row++) {
    data.set(
      mapped.subarray(row * bytesPerRow, row * bytesPerRow + rowBytes),
      row * rowBytes,
    );
  }
  readback.unmap();
  untrackResource(ctx.device, "buffers", readbackId);
  readback.destroy();
  return data;
}

/**
 * Asynchronously read a buffer back to the CPU.
 *
 * ```js
 * const data = new Float32Array(await gpu.readBuffer(ctx, particleBuffer));
 * ```
 */
export async function readBuffer(
  ctx: GpuContext,
  source: GpuBuffer,
  byteOffset = 0,
  byteLength = source.size - byteOffset,
): Promise<ArrayBuffer> {
  // Copy sizes must be multiples of 4 — copy the aligned range (buffer sizes
  // are 4-aligned at creation, so it cannot overrun) and slice the result
  const copyLength = alignTo(byteLength, 4);
  const readback = ctx.device.createBuffer({
    label: "pex-gpu readback",
    size: copyLength,
    usage: BUFFER_USAGE_PRESETS.readback,
  });
  const readbackId = nextId();
  trackResource(ctx.device, "buffers", {
    id: readbackId,
    label: readback.label,
    usage: readback.usage,
    bytes: readback.size,
  });

  const encoder = ctx.device.createCommandEncoder({
    label: "pex-gpu readBuffer",
  });
  encoder.copyBufferToBuffer(
    source.buffer,
    byteOffset,
    readback,
    0,
    copyLength,
  );
  ctx.device.queue.submit([encoder.finish()]);

  await readback.mapAsync(GPUMapMode.READ);
  const data = readback.getMappedRange().slice(0, byteLength);
  readback.unmap();
  untrackResource(ctx.device, "buffers", readbackId);
  readback.destroy();
  return data;
}