import type { BufferData, TypedArray } from "./types.js";

/** Round `value` up to the next multiple of `alignment`. */
export const alignTo = (value: number, alignment: number): number =>
  Math.ceil(value / alignment) * alignment;

// ArrayBuffer.isView() is also true for DataView, which has no element type
// or .length — excluded here since it isn't a typed array.
export const isTypedArray = (data: unknown): data is TypedArray =>
  ArrayBuffer.isView(data) && !(data instanceof DataView);

/**
 * Coerce buffer factory input into a typed array. Nested number arrays are
 * flattened; plain arrays default to Float32Array (Uint16Array/Uint32Array for
 * index data, based on the max value). Raw ArrayBuffers pass through as a byte
 * view (Uint8Array) — the data is trusted to already be in the right format.
 */
export function toTypedArray(
  data: BufferData,
  { index = false }: { index?: boolean } = {},
): TypedArray {
  if (isTypedArray(data)) return data;
  if (data instanceof ArrayBuffer) return new Uint8Array(data);

  const flat: number[] = Array.isArray(data[0])
    ? (data as number[][]).flat()
    : (data as number[]);

  if (index) {
    return flat.some((value) => value > 0xff_ff)
      ? new Uint32Array(flat)
      : new Uint16Array(flat);
  }
  return new Float32Array(flat);
}

// Half-float formats: the only float formats whose storage differs from the
// Float32Array default, so a plain number array must be packed explicitly.
const FLOAT16_FORMATS = new Set<GPUTextureFormat>([
  "r16float",
  "rg16float",
  "rgba16float",
]);

/**
 * Coerce texture upload data to a typed array. Typed arrays pass through as
 * trusted bytes — the required path for integer/unorm/snorm/packed formats. A
 * plain number array is treated as float data and packed at the format's float
 * precision: Float16Array for half-float formats, Float32Array otherwise.
 */
export function toTextureData(
  data: BufferData,
  format: GPUTextureFormat,
): TypedArray {
  if (Array.isArray(data) && FLOAT16_FORMATS.has(format)) {
    return new Float16Array(
      Array.isArray(data[0]) ? (data as number[][]).flat() : (data as number[]),
    );
  }
  return toTypedArray(data);
}

/** Bytes per texel for formats supported by updateTexture/readTexture. */
const BYTES_PER_TEXEL: Partial<Record<GPUTextureFormat, number>> = {
  r8unorm: 1,
  r8snorm: 1,
  r8uint: 1,
  r8sint: 1,
  rg8unorm: 2,
  rg8snorm: 2,
  rg8uint: 2,
  rg8sint: 2,
  r16float: 2,
  r16uint: 2,
  r16sint: 2,
  rg16float: 4,
  rg16uint: 4,
  rg16sint: 4,
  r32float: 4,
  r32uint: 4,
  r32sint: 4,
  rgba8unorm: 4,
  "rgba8unorm-srgb": 4,
  rgba8snorm: 4,
  rgba8uint: 4,
  rgba8sint: 4,
  bgra8unorm: 4,
  "bgra8unorm-srgb": 4,
  rgb9e5ufloat: 4,
  rgb10a2uint: 4,
  rgb10a2unorm: 4,
  rg11b10ufloat: 4,
  rg32float: 8,
  rg32uint: 8,
  rg32sint: 8,
  rgba16float: 8,
  rgba16uint: 8,
  rgba16sint: 8,
  rgba32float: 16,
  rgba32uint: 16,
  rgba32sint: 16,
};

/**
 * Bytes per texel for a supported texture format. Throws for
 * unsupported/compressed formats — compressed data is measured in blocks, see
 * {@link blockInfo}.
 */
export function bytesPerTexel(format: GPUTextureFormat): number {
  const bytes = BYTES_PER_TEXEL[format];
  if (!bytes) {
    throw new Error(`pex-gpu: unsupported texture data format "${format}"`);
  }
  return bytes;
}

/** Footprint of one compressed block: its texel extent and encoded size. */
export interface BlockInfo {
  width: number;
  height: number;
  bytes: number;
}

// Bytes per 4x4 block for the fixed-size families. ASTC is omitted: every ASTC
// block is 16 bytes and its extent is in the format name, so it is parsed
// rather than enumerated (52 compressed formats, 28 of them ASTC).
const BLOCK_BYTES: Partial<Record<GPUTextureFormat, number>> = {
  "bc1-rgba-unorm": 8,
  "bc1-rgba-unorm-srgb": 8,
  "bc2-rgba-unorm": 16,
  "bc2-rgba-unorm-srgb": 16,
  "bc3-rgba-unorm": 16,
  "bc3-rgba-unorm-srgb": 16,
  "bc4-r-unorm": 8,
  "bc4-r-snorm": 8,
  "bc5-rg-unorm": 16,
  "bc5-rg-snorm": 16,
  "bc6h-rgb-ufloat": 16,
  "bc6h-rgb-float": 16,
  "bc7-rgba-unorm": 16,
  "bc7-rgba-unorm-srgb": 16,
  "etc2-rgb8unorm": 8,
  "etc2-rgb8unorm-srgb": 8,
  "etc2-rgb8a1unorm": 8,
  "etc2-rgb8a1unorm-srgb": 8,
  "etc2-rgba8unorm": 16,
  "etc2-rgba8unorm-srgb": 16,
  "eac-r11unorm": 8,
  "eac-r11snorm": 8,
  "eac-rg11unorm": 16,
  "eac-rg11snorm": 16,
};

const ASTC_BLOCK = /^astc-(\d+)x(\d+)-unorm(-srgb)?$/;

/**
 * Block footprint of a compressed format, or `undefined` for uncompressed ones
 * — which is also the test for "is this format block-compressed".
 */
export function blockInfo(format: GPUTextureFormat): BlockInfo | undefined {
  const bytes = BLOCK_BYTES[format];
  if (bytes) return { width: 4, height: 4, bytes };

  const astc = ASTC_BLOCK.exec(format);
  return astc
    ? { width: Number(astc[1]), height: Number(astc[2]), bytes: 16 }
    : undefined;
}

/** Whether `format` stores texels in compressed blocks. */
export const isCompressedFormat = (format: GPUTextureFormat): boolean =>
  blockInfo(format) !== undefined;

/**
 * Row pitch for a `writeTexture` of `width`x`height` texels. Compressed formats
 * are measured in whole blocks, partial blocks at the edge counting as full
 * ones. `queue.writeTexture` imposes no 256-byte row alignment (that is a
 * `copyBufferToTexture` rule), so this is the tight pitch.
 */
export function texelCopyLayout(
  format: GPUTextureFormat,
  width: number,
  height: number,
): { bytesPerRow: number; rowsPerImage: number } {
  const block = blockInfo(format);
  if (!block) {
    return { bytesPerRow: bytesPerTexel(format) * width, rowsPerImage: height };
  }
  return {
    bytesPerRow: Math.ceil(width / block.width) * block.bytes,
    rowsPerImage: Math.ceil(height / block.height),
  };
}

/**
 * Copy extent rounded up to whole blocks — the physical size a compressed
 * subresource occupies, which copy extents must match even when the logical mip
 * size is smaller. Identity for uncompressed formats.
 */
export function physicalExtent(
  format: GPUTextureFormat,
  width: number,
  height: number,
): { width: number; height: number } {
  const block = blockInfo(format);
  if (!block) return { width, height };
  return {
    width: Math.ceil(width / block.width) * block.width,
    height: Math.ceil(height / block.height) * block.height,
  };
}

/** Compressed textures are allocated in whole blocks. */
export function assertBlockAlignedSize(
  format: GPUTextureFormat,
  width: number,
  height: number,
): void {
  const block = blockInfo(format);
  if (!block) return;
  if (width % block.width || height % block.height) {
    throw new Error(
      `pex-gpu: "${format}" requires ${block.width}×${block.height} block-aligned dimensions, got ${width}×${height}`,
    );
  }
}

/** Compressed copies must start on a block boundary. */
export function assertBlockAlignedOrigin(
  format: GPUTextureFormat,
  origin: GPUOrigin3D,
): void {
  const block = blockInfo(format);
  if (!block) return;
  const [x = 0, y = 0] = Array.isArray(origin)
    ? origin
    : [origin.x ?? 0, origin.y ?? 0];
  if (x % block.width || y % block.height) {
    throw new Error(
      `pex-gpu: "${format}" copies must start on a ${block.width}×${block.height} block boundary, got origin ${x},${y}`,
    );
  }
}