import { alignTo, isTypedArray } from "./utils.js";
import type { StructLayout, StructMember, TypedArray } from "./types.js";

/** A value writable into a uniform/storage struct member. */
export type PackableValue =
  number | boolean | number[] | number[][] | TypedArray | PackableStruct;

/** A member that is itself a struct: an object keyed by its member names. */
export interface PackableStruct {
  [member: string]: PackableValue;
}

/** Written straight through; everything else is flattened first. */
const isScalar = (value: PackableValue): value is number | boolean =>
  typeof value === "number" || typeof value === "boolean";

const toFlatArray = (
  value: Exclude<PackableValue, number | boolean>,
): ArrayLike<number> =>
  Array.isArray(value) && Array.isArray(value[0])
    ? (value as number[][]).flat()
    : (value as ArrayLike<number>);

// Matches i32/u32 scalars or vecN<i32/u32> — selects int vs. float writes below.
const isIntegerType = (type: string) => /(^|<)[iu]32>?$/.test(type);

function writeScalar(
  view: DataView,
  offset: number,
  type: string,
  value: number,
) {
  if (type.includes("u32")) view.setUint32(offset, value, true);
  else if (type.includes("i32")) view.setInt32(offset, value, true);
  else view.setFloat32(offset, value, true);
}

/**
 * Write one value at a byte offset, honoring WGSL layout rules. Matrices are
 * written column by column with the column stride of their vector type —
 * mat3x3<f32> columns are 16-byte aligned, not 12.
 */
function writeValue(
  view: DataView,
  offset: number,
  type: string,
  value: PackableValue,
) {
  if (isScalar(value)) {
    writeScalar(view, offset, type, Number(value));
    return;
  }

  const data = toFlatArray(value);
  const matrix = /^mat([234])x([234])</.exec(type);
  if (matrix) {
    const [cols, rows] = [Number(matrix[1]), Number(matrix[2])];
    // Uniform address space: vec2 columns align to 8 bytes, vec3/vec4
    // columns align to 16 (mat3x3's 12-byte columns are the classic trap).
    const columnStride = alignTo(rows * 4, rows === 2 ? 8 : 16);
    for (let col = 0; col < cols; col++) {
      for (let row = 0; row < rows; row++) {
        view.setFloat32(
          offset + col * columnStride + row * 4,
          data[col * rows + row]!,
          true,
        );
      }
    }
    return;
  }

  const integer = isIntegerType(type);
  for (let i = 0; i < data.length; i++) {
    if (integer) writeScalar(view, offset + i * 4, type, data[i]!);
    else view.setFloat32(offset + i * 4, data[i]!, true);
  }
}

function writeMember(
  view: DataView,
  member: StructMember,
  value: PackableValue,
) {
  if (
    member.elementLayout &&
    member.arrayCount !== undefined &&
    member.arrayStride !== undefined
  ) {
    // Array of structs: pack each element by member name at the array stride.
    const elements = value as unknown as Record<string, PackableValue>[];
    for (let i = 0; i < member.arrayCount; i++) {
      if (elements[i] === undefined) continue;
      packStruct(
        view.buffer as ArrayBuffer,
        member.elementLayout,
        elements[i]!,
        view.byteOffset + member.offset + i * member.arrayStride,
      );
    }
    return;
  }
  if (member.arrayCount !== undefined && member.arrayStride !== undefined) {
    // Fixed-size array: elements are placed at the (possibly padded) stride.
    const elementType =
      /^array<(.+?)(?:,\s*\d+\s*)?>$/.exec(member.type)?.[1] ?? "f32";
    const data = toFlatArray(isScalar(value) ? [Number(value)] : value);
    const componentsPerElement = data.length / member.arrayCount;
    for (let i = 0; i < member.arrayCount; i++) {
      const element = Array.prototype.slice.call(
        data,
        i * componentsPerElement,
        (i + 1) * componentsPerElement,
      );
      writeValue(
        view,
        member.offset + i * member.arrayStride,
        elementType,
        element,
      );
    }
    return;
  }
  if (member.structLayout) {
    packStruct(
      view.buffer as ArrayBuffer,
      member.structLayout,
      value as PackableStruct,
      view.byteOffset + member.offset,
    );
    return;
  }
  writeValue(view, member.offset, member.type, value);
}

/**
 * Values reach here from an object the caller assembled, so a member holding an
 * `undefined` — a default that was never applied, a property read off the wrong
 * object — is the common mistake. Left to the writers it surfaces as a `length`
 * of undefined inside a DataView write, naming neither the struct nor the
 * member; a value of the wrong shape is worse still, writing nothing at all.
 */
function assertPackable(
  layout: StructLayout,
  member: StructMember,
  value: PackableValue,
): void {
  const label = member.name
    ? `member "${member.name}" of struct "${layout.name}" (${member.type})`
    : `binding of type "${member.type}"`;

  if (value === undefined || value === null) {
    throw new Error(
      `pex-gpu: ${label} is ${value}. Omit the key to leave the member unwritten.`,
    );
  }
  if (member.structLayout) {
    if (isScalar(value) || Array.isArray(value) || isTypedArray(value)) {
      throw new Error(
        `pex-gpu: ${label} is a struct, expected an object of its members`,
      );
    }
    return; // Its own members are checked as they are packed.
  }
  if (!isScalar(value) && !Array.isArray(value) && !isTypedArray(value)) {
    throw new Error(
      `pex-gpu: ${label} expects a number, boolean, array or typed array, got ${typeof value}`,
    );
  }
}

/**
 * Pack values into a struct layout inside `destination`.
 *
 * `values` is either an object keyed by member name, or — when the layout is a
 * synthetic single-value layout for a non-struct binding — the value itself.
 * Throws on unknown member names, and on values that cannot be written, to
 * surface typos and missing defaults early.
 *
 * Known limitation: elements of `array<SomeStruct, N>` members are written
 * tightly packed — structs with internal padding (eg. two vec3f members) need
 * their padding included in the source data.
 */
export function packStruct(
  destination: ArrayBuffer,
  layout: StructLayout,
  values: Record<string, PackableValue> | PackableValue,
  byteOffset = 0,
): void {
  const view = new DataView(destination, byteOffset);

  const single = layout.members.length === 1 && layout.members[0]!.name === "";
  if (single) {
    assertPackable(layout, layout.members[0]!, values as PackableValue);
    writeMember(view, layout.members[0]!, values as PackableValue);
    return;
  }

  for (const [name, value] of Object.entries(
    values as Record<string, PackableValue>,
  )) {
    const member = layout.members.find((candidate) => candidate.name === name);
    if (!member) {
      throw new Error(
        `pex-gpu: unknown member "${name}" in struct "${layout.name}" (members: ${layout.members.map((m) => m.name).join(", ")})`,
      );
    }
    assertPackable(layout, member, value);
    writeMember(view, member, value);
  }
}