import { isGpuBuffer } from "./buffer.js";
import type {
  Attributes,
  GpuBuffer,
  VertexAttribute,
  VertexInput,
} from "./types.js";

export interface ResolvedVertexState {
  layouts: GPUVertexBufferLayout[];
  /** Buffer bound to each slot, matching layouts order. */
  bindings: { buffer: GpuBuffer; offset: number }[];
  /** Cache key participating in pipeline lookup. */
  key: string;
  /** Vertex count inferred from the first attribute, when available. */
  inferredCount?: number;
}

// e.g. "float32x3" → 3; formats without an x-suffix (eg. "float32") are single-component.
const componentCount = (format: GPUVertexFormat): number =>
  Number(/x(\d)$/.exec(format)?.[1] ?? 1);

/**
 * Derive vertex buffer layouts from the shader's reflected inputs and the
 * command's attributes: one buffer slot per attribute, formats and strides
 * inferred from WGSL types unless overridden.
 */
export function resolveVertexState(
  vertexInputs: readonly VertexInput[],
  attributes: Attributes,
): ResolvedVertexState {
  const layouts: GPUVertexBufferLayout[] = [];
  const bindings: ResolvedVertexState["bindings"] = [];
  const keyParts: string[] = [];
  let inferredCount: number | undefined;

  for (const input of vertexInputs) {
    const value = attributes[input.name];
    if (!value) {
      throw new Error(
        `pex-gpu: missing attribute "${input.name}" (@location(${input.location})); provided: ${Object.keys(attributes).join(", ") || "none"}`,
      );
    }

    const attribute: VertexAttribute = isGpuBuffer(value)
      ? { buffer: value }
      : value;
    const format = attribute.format ?? input.format;
    const arrayStride = attribute.arrayStride ?? input.arrayStride;
    const stepMode = attribute.stepMode ?? "vertex";

    layouts.push({
      arrayStride,
      stepMode,
      attributes: [{ shaderLocation: input.location, offset: 0, format }],
    });
    bindings.push({ buffer: attribute.buffer, offset: attribute.offset ?? 0 });
    keyParts.push(`${input.location}:${format}:${arrayStride}:${stepMode}`);

    // Only infer from a plain per-vertex attribute with a known element
    // count and no custom stride — instance-stepped or custom-strided
    // (eg. interleaved) buffers can't reliably derive a vertex count this way.
    if (
      inferredCount === undefined &&
      stepMode === "vertex" &&
      attribute.buffer.length !== undefined &&
      attribute.arrayStride === undefined
    ) {
      inferredCount = attribute.buffer.length / componentCount(format);
    }
  }

  return {
    layouts,
    bindings,
    key: keyParts.join("|"),
    ...(inferredCount !== undefined && { inferredCount }),
  };
}