import { alignTo } from "./utils.js";
import { debugLog, recordCacheAccess } from "./debug.js";
import { shared } from "./shared.js";
import type {
  BindingKind,
  ShaderBinding,
  ShaderReflection,
  StructLayout,
  StructMember,
  VertexInput,
} from "./types.js";

/**
 * Hand-written WGSL reflection.
 *
 * Extracts, from a WGSL source string:
 *
 * - Struct definitions with byte layouts (uniform address space rules)
 * - Resource bindings (`@group`/`@binding` var declarations)
 * - Vertex inputs (`@location` parameters of the `@vertex` entry point)
 * - Entry point names per stage
 *
 * This is intentionally not a full WGSL parser: it covers the declaration
 * grammar needed to build bind group layouts and pack uniform data. Shaders
 * using exotic constructs can bypass reflection with explicit bind groups.
 */

// ─── Scalar/vector/matrix type info ──────────────────────────────────────────

interface TypeInfo {
  size: number;
  align: number;
  vertexFormat?: GPUVertexFormat;
}

const SCALARS: Record<string, TypeInfo & { suffix: string }> = {
  f32: { size: 4, align: 4, suffix: "float32", vertexFormat: "float32" },
  i32: { size: 4, align: 4, suffix: "sint32", vertexFormat: "sint32" },
  u32: { size: 4, align: 4, suffix: "uint32", vertexFormat: "uint32" },
  f16: { size: 2, align: 2, suffix: "float16" },
};

/** Normalize predeclared aliases: vec3f → vec3<f32>, mat4x4f → mat4x4<f32>, … */
export function normalizeType(type: string): string {
  return type
    .trim()
    .replaceAll(/\s+/g, "")
    .replaceAll(
      /\b(vec[234]|mat[234]x[234])([fiuh])\b/g,
      (_, base: string, suffix: string) =>
        `${base}<${{ f: "f32", i: "i32", u: "u32", h: "f16" }[suffix]}>`,
    );
}

/**
 * Which layout rules a struct is measured under. The uniform address space
 * rounds array strides and nested struct alignments up to 16; storage uses
 * natural alignment throughout. The same struct bound both ways genuinely has
 * two different layouts, which is why every layout here is keyed by space as
 * well as by name — packing a storage struct under the uniform rules writes
 * every member at the wrong offset, and nothing downstream can detect it.
 */
export type LayoutSpace = "uniform" | "storage";

export const layoutKey = (space: LayoutSpace, name: string) =>
  `${space}:${name}`;

function typeInfo(
  type: string,
  space: LayoutSpace,
  structs: Map<string, StructLayout>,
): TypeInfo & { structLayout?: StructLayout } {
  const scalar = SCALARS[type];
  if (scalar) return scalar;

  let match = /^vec([234])<(\w+)>$/.exec(type);
  if (match) {
    const n = Number(match[1]);
    const component = SCALARS[match[2]!];
    if (!component)
      throw reflectionError(`unknown vector component in "${type}"`);
    const vertexFormat =
      component.suffix === "float16" && n === 3
        ? undefined // no float16x3 vertex format exists
        : (`${component.suffix}x${n}` as GPUVertexFormat);
    return {
      size: n * component.size,
      align: (n === 3 ? 4 : n) * component.align,
      ...(vertexFormat && { vertexFormat }),
    };
  }

  match = /^mat([234])x([234])<(\w+)>$/.exec(type);
  if (match) {
    const [cols, rows] = [Number(match[1]), Number(match[2])];
    const column = typeInfo(`vec${rows}<${match[3]}>`, space, structs);
    const stride = alignTo(column.size, column.align);
    return { size: cols * stride, align: column.align };
  }

  const struct = structs.get(layoutKey(space, type));
  if (struct)
    return { size: struct.size, align: struct.align, structLayout: struct };

  throw reflectionError(`unsupported type "${type}" in buffer layout`);
}

/** Element stride and alignment of an array<T, N> under `space`'s rules. */
function arrayInfo(
  elementType: string,
  space: LayoutSpace,
  structs: Map<string, StructLayout>,
) {
  const element = typeInfo(elementType, space, structs);
  const stride = alignTo(element.size, element.align);
  return space === "uniform"
    ? { stride: alignTo(stride, 16), align: alignTo(element.align, 16) }
    : { stride, align: element.align };
}

const reflectionError = (message: string) =>
  new Error(`pex-gpu: WGSL reflection: ${message}`);

// ─── Source preprocessing ────────────────────────────────────────────────────

function stripComments(source: string): string {
  let out = "";
  let i = 0;
  while (i < source.length) {
    if (source.startsWith("//", i)) {
      const end = source.indexOf("\n", i);
      i = end === -1 ? source.length : end;
    } else if (source.startsWith("/*", i)) {
      // WGSL block comments nest (unlike C/JS), so track depth rather than
      // stopping at the first closing */.
      let depth = 1;
      i += 2;
      while (i < source.length && depth > 0) {
        if (source.startsWith("/*", i)) (depth++, (i += 2));
        else if (source.startsWith("*/", i)) (depth--, (i += 2));
        else i++;
      }
    } else {
      out += source[i++];
    }
  }
  return out;
}

/**
 * Split on commas at nesting depth 0. Defaults suit type expressions; argument
 * lists pass `"(["`/`")]"` instead, where `<` and `>` are comparisons.
 */
function splitTopLevel(input: string, open = "<(", close = ">)"): string[] {
  const parts: string[] = [];
  let depth = 0;
  let current = "";
  for (const char of input) {
    if (open.includes(char)) depth++;
    else if (close.includes(char)) depth--;
    if (char === "," && depth === 0) {
      parts.push(current);
      current = "";
    } else {
      current += char;
    }
  }
  if (current.trim()) parts.push(current);
  return parts;
}

interface Attributes {
  location?: number;
  align?: number;
  size?: number;
  builtin?: string;
  group?: number;
  binding?: number;
}

function parseAttributes(input: string): Attributes {
  const attributes: Attributes = {};
  for (const [, name, value] of input.matchAll(
    /@(\w+)\s*\(\s*([^)]*?)\s*\)/g,
  )) {
    if (name === "builtin") attributes.builtin = value!;
    else if (
      name === "location" ||
      name === "align" ||
      name === "size" ||
      name === "group" ||
      name === "binding"
    ) {
      attributes[name] = Number(value);
    }
  }
  return attributes;
}

// ─── Struct parsing ──────────────────────────────────────────────────────────

interface RawMember {
  name: string;
  type: string;
  attributes: Attributes;
}

function parseStructs(source: string): Map<string, RawMember[]> {
  const structs = new Map<string, RawMember[]>();
  for (const [, name, body] of source.matchAll(
    /struct\s+(\w+)\s*\{([^}]*)\}/g,
  )) {
    const members: RawMember[] = [];
    for (const entry of splitTopLevel(body!)) {
      const match = /^((?:\s*@\w+\s*\([^)]*\)\s*)*)\s*(\w+)\s*:\s*(.+)$/s.exec(
        entry,
      );
      if (!match) continue;
      members.push({
        name: match[2]!,
        type: normalizeType(match[3]!),
        attributes: parseAttributes(match[1]!),
      });
    }
    structs.set(name!, members);
  }
  return structs;
}

/**
 * Compute (and cache) a struct's byte layout under `space`'s rules, recursing
 * into nested struct and array element types.
 */
function computeLayout(
  name: string,
  space: LayoutSpace,
  rawStructs: Map<string, RawMember[]>,
  layouts: Map<string, StructLayout>,
): StructLayout {
  const cached = layouts.get(layoutKey(space, name));
  if (cached) return cached;

  const raw = rawStructs.get(name);
  if (!raw) throw reflectionError(`unknown struct "${name}"`);

  let offset = 0;
  let structAlign = 0;
  const members: StructMember[] = [];

  for (const { name: memberName, type, attributes } of raw) {
    if (attributes.builtin !== undefined || attributes.location !== undefined) {
      continue; // IO struct member, not host-shareable data
    }

    let info: TypeInfo & { structLayout?: StructLayout };
    let arrayCount: number | undefined;
    let arrayStride: number | undefined;

    const array = /^array<(.+?)(?:,\s*(\d+)\s*)?>$/.exec(type);
    if (array) {
      const elementType = array[1]!;
      if (rawStructs.has(elementType))
        computeLayout(elementType, space, rawStructs, layouts);
      const element = arrayInfo(elementType, space, layouts);
      arrayCount = array[2] ? Number(array[2]) : undefined;
      arrayStride = element.stride;
      if (arrayCount === undefined) {
        throw reflectionError(
          `runtime-sized array "${memberName}" is only allowed as the last member of a storage buffer struct; bind it as a whole buffer instead`,
        );
      }
      info = { size: arrayCount * element.stride, align: element.align };
    } else {
      if (rawStructs.has(type)) computeLayout(type, space, rawStructs, layouts);
      info = typeInfo(type, space, layouts);
      // Uniform address space: nested struct members align to 16. Only the
      // strict reading of the rule — Dawn uses the struct's natural alignment
      // — so a nested struct member under 16-byte alignment wants an explicit
      // `@align` in the WGSL rather than either of us guessing.
      if (space === "uniform" && info.structLayout) {
        info = { ...info, align: alignTo(info.align, 16) };
      }
    }

    const align = attributes.align ?? info.align;
    offset = alignTo(offset, align);
    members.push({
      name: memberName,
      type,
      offset,
      size: attributes.size ?? info.size,
      align,
      ...(arrayCount !== undefined && { arrayCount }),
      ...(arrayStride !== undefined && { arrayStride }),
      ...(info.structLayout && { structLayout: info.structLayout }),
    });
    offset += attributes.size ?? info.size;
    structAlign = Math.max(structAlign, align);
  }

  const layout: StructLayout = {
    name,
    align: structAlign || 4,
    size: alignTo(offset, structAlign || 4),
    members,
  };
  layouts.set(layoutKey(space, name), layout);
  return layout;
}

// ─── Binding parsing ─────────────────────────────────────────────────────────

const TEXTURE_TYPES: Record<
  string,
  {
    viewDimension: GPUTextureViewDimension;
    depth?: boolean;
    multisampled?: boolean;
  }
> = {
  texture_1d: { viewDimension: "1d" },
  texture_2d: { viewDimension: "2d" },
  texture_2d_array: { viewDimension: "2d-array" },
  texture_3d: { viewDimension: "3d" },
  texture_cube: { viewDimension: "cube" },
  texture_cube_array: { viewDimension: "cube-array" },
  texture_multisampled_2d: { viewDimension: "2d", multisampled: true },
  texture_depth_2d: { viewDimension: "2d", depth: true },
  texture_depth_2d_array: { viewDimension: "2d-array", depth: true },
  texture_depth_cube: { viewDimension: "cube", depth: true },
  texture_depth_cube_array: { viewDimension: "cube-array", depth: true },
  texture_depth_multisampled_2d: {
    viewDimension: "2d",
    depth: true,
    multisampled: true,
  },
};

const SAMPLE_TYPES: Record<string, GPUTextureSampleType> = {
  f32: "float",
  u32: "uint",
  i32: "sint",
};

const STORAGE_TEXTURE_ACCESS: Record<string, GPUStorageTextureAccess> = {
  read: "read-only",
  write: "write-only",
  read_write: "read-write",
};

function parseBindings(
  source: string,
  structs: Map<string, StructLayout>,
  rawStructs: Map<string, RawMember[]>,
  visibility: GPUShaderStageFlags,
): ShaderBinding[] {
  const bindings: ShaderBinding[] = [];
  const declarations = source.matchAll(
    /((?:@\w+\s*\([^)]*\)\s*)+)var\s*(?:<([^>]*)>)?\s+(\w+)\s*:\s*([^;=]+)/g,
  );

  for (const [, attributeText, addressSpace, name, rawType] of declarations) {
    const { group, binding } = parseAttributes(attributeText!);
    if (group === undefined || binding === undefined) continue;

    const type = normalizeType(rawType!);
    const base: Pick<ShaderBinding, "name" | "group" | "binding"> = {
      name: name!,
      group,
      binding,
    };

    if (addressSpace?.startsWith("uniform")) {
      const layout = resolveBufferLayout(type, "uniform", structs, rawStructs);
      bindings.push({
        ...base,
        kind: "uniform",
        visibility,
        structLayout: layout,
        minBindingSize: layout.size,
      });
    } else if (addressSpace?.startsWith("storage")) {
      const readOnly = !addressSpace.includes("read_write");
      const kind: BindingKind = readOnly ? "read-only-storage" : "storage";
      // A runtime-sized array of a known struct can be filled from a JS array.
      const elementType = /^array<([^,]+)>$/.exec(type)?.[1];
      const element =
        elementType && rawStructs.has(elementType)
          ? computeLayout(elementType, "storage", rawStructs, structs)
          : undefined;
      bindings.push({
        ...base,
        kind,
        ...(element && {
          arrayElementLayout: element,
          arrayStride: alignTo(element.size, element.align),
          // One element is the smallest valid binding — and it is what
          // distinguishes two layouts that are otherwise identical, so it has
          // to live on the binding where the layout cache key can see it.
          minBindingSize: alignTo(element.size, element.align),
        }),
        // read_write storage is not allowed in the vertex stage
        visibility: readOnly ? visibility : visibility & ~GPUShaderStage.VERTEX,
      });
    } else if (type === "sampler" || type === "sampler_comparison") {
      bindings.push({
        ...base,
        kind: type === "sampler" ? "sampler" : "comparison-sampler",
        visibility,
      });
    } else if (type === "texture_external") {
      bindings.push({ ...base, kind: "external-texture", visibility });
    } else {
      const textureMatch = /^(\w+?)(?:<(.+)>)?$/.exec(type);
      const textureType = textureMatch && TEXTURE_TYPES[textureMatch[1]!];
      if (textureType) {
        bindings.push({
          ...base,
          kind: "texture",
          visibility,
          viewDimension: textureType.viewDimension,
          sampleType: textureType.depth
            ? "depth"
            : (SAMPLE_TYPES[textureMatch![2] ?? "f32"] ?? "float"),
          ...(textureType.multisampled && { multisampled: true }),
        });
      } else if (
        textureMatch &&
        textureMatch[1]!.startsWith("texture_storage_")
      ) {
        const [format, access = "write"] = (textureMatch[2] ?? "").split(
          ",",
          2,
        );
        bindings.push({
          ...base,
          kind: "storage-texture",
          visibility:
            access.trim() === "read"
              ? visibility
              : // Mirrors the read_write-storage-buffer rule above: write
                // access is not allowed in the vertex stage.
                visibility & ~GPUShaderStage.VERTEX,
          viewDimension:
            TEXTURE_TYPES[textureMatch[1]!.replace("_storage", "")]
              ?.viewDimension ?? "2d",
          storageTextureFormat: format!.trim() as GPUTextureFormat,
          storageTextureAccess:
            STORAGE_TEXTURE_ACCESS[access.trim()] ?? "write-only",
        });
      } else {
        throw reflectionError(
          `unsupported binding type "${type}" for "${name}"`,
        );
      }
    }
  }
  return bindings;
}

/** Layout for a buffer binding: struct type, or a single-value synthetic layout. */
function resolveBufferLayout(
  type: string,
  space: LayoutSpace,
  structs: Map<string, StructLayout>,
  rawStructs: Map<string, RawMember[]>,
): StructLayout {
  if (rawStructs.has(type))
    return computeLayout(type, space, rawStructs, structs);

  // array<T, N> binding, bound directly (not wrapped in a struct) — e.g. a
  // joint-matrix uniform. Struct element types (light arrays) also carry the
  // element layout so each element is packed by member name.
  const array = /^array<(.+?)(?:,\s*(\d+)\s*)?>$/.exec(type);
  if (array && array[2]) {
    const elementType = array[1]!;
    const count = Number(array[2]);
    const isStructElement = rawStructs.has(elementType);
    if (isStructElement) computeLayout(elementType, space, rawStructs, structs);
    const { stride, align } = arrayInfo(elementType, space, structs);
    const size = stride * count;
    return {
      name: "",
      size,
      align,
      members: [
        {
          name: "",
          type,
          offset: 0,
          size,
          align,
          arrayCount: count,
          arrayStride: stride,
          ...(isStructElement && {
            elementLayout: structs.get(layoutKey(space, elementType))!,
          }),
        },
      ],
    };
  }

  const info = typeInfo(type, space, structs);
  return {
    name: "",
    size: alignTo(info.size, info.align),
    align: info.align,
    members: [
      { name: "", type, offset: 0, size: info.size, align: info.align },
    ],
  };
}

// ─── Entry points and vertex inputs ──────────────────────────────────────────

function parseEntryPoints(source: string) {
  const entryPoints: { vertex?: string; fragment?: string; compute?: string } =
    {};
  const vertexParams: string[] = [];

  const stageAttribute =
    /@(vertex|fragment|compute)\s*(?:@\w+\s*\([^)]*\)\s*)*fn\s+(\w+)\s*\(/g;
  for (const match of source.matchAll(stageAttribute)) {
    const stage = match[1] as "vertex" | "fragment" | "compute";
    entryPoints[stage] ??= match[2]!;
    if (stage === "vertex" && entryPoints.vertex === match[2]) {
      // Balanced scan of the parameter list (attributes contain parentheses).
      let depth = 1;
      let i = match.index! + match[0].length;
      const start = i;
      while (i < source.length && depth > 0) {
        if (source[i] === "(") depth++;
        else if (source[i] === ")") depth--;
        i++;
      }
      vertexParams.push(source.slice(start, i - 1));
    }
  }
  return { entryPoints, vertexParams };
}

function parseVertexInputs(
  paramsText: string,
  rawStructs: Map<string, RawMember[]>,
): VertexInput[] {
  const inputs: VertexInput[] = [];

  const addInput = (name: string, type: string, location: number) => {
    const info = typeInfo(type, "storage", new Map());
    if (!info.vertexFormat) {
      throw reflectionError(
        `type "${type}" of vertex input "${name}" has no vertex format`,
      );
    }
    inputs.push({
      name,
      location,
      type,
      format: info.vertexFormat,
      arrayStride: info.size,
    });
  };

  for (const param of splitTopLevel(paramsText)) {
    const match = /^((?:\s*@\w+\s*\([^)]*\)\s*)*)\s*(\w+)\s*:\s*(.+)$/s.exec(
      param,
    );
    if (!match) continue;
    const attributes = parseAttributes(match[1]!);
    if (attributes.builtin !== undefined) continue;
    const type = normalizeType(match[3]!);

    if (attributes.location === undefined) {
      // Struct-typed parameter: expand @location members.
      for (const member of rawStructs.get(type) ?? []) {
        if (member.attributes.location === undefined) continue;
        addInput(member.name, member.type, member.attributes.location);
      }
    } else {
      addInput(match[2]!, type, attributes.location);
    }
  }
  return inputs.sort((a, b) => a.location - b.location);
}

// ─── Override constants ──────────────────────────────────────────────────────

/** Per-function: the traced names it references and the functions it calls. */
interface FunctionNode {
  uses: Set<string>;
  calls: Set<string>;
}

interface ReachabilityAnalysis {
  names: Set<string>;
  graph: Map<string, FunctionNode>;
  /**
   * Names referenced outside any function (array sizes, `@workgroup_size`,
   * override-to-override defaults) — these apply to every entry point.
   */
  moduleScopeUsed: Set<string>;
}

/** An override declaration, so its own name doesn't read as a use. */
const OVERRIDE_DECLARATION = /\boverride\s+[A-Za-z_]\w*[^;]*;/g;
/** Likewise for a binding: ``@group`(n) `@binding`(n) var<space> name: type;`. */
const BINDING_DECLARATION =
  /@group\s*\([^)]*\)\s*@binding\s*\([^)]*\)\s*var(?:<[^>]*>)?\s+[A-Za-z_]\w*\s*:[^;]*;/g;

// `code` is already comment-stripped (parseWGSL shares its stripped source), so
// an identifier scan is unambiguous — WGSL has no string literals. WGSL has no
// indirection through bindings either, so every use is a literal occurrence of
// the name and the trace cannot miss one.
function analyzeReachability(
  code: string,
  names: Set<string>,
  declarations: RegExp,
): ReachabilityAnalysis {
  // Nothing to trace: skip the (comparatively costly) function scan entirely.
  if (!names.size)
    return { names, graph: new Map(), moduleScopeUsed: new Set() };

  // Pass 1: function name → body identifiers (a callee may be defined later, so
  // the call graph can only be reduced once every name is known).
  const bodies = new Map<string, string[]>();
  const spans: [number, number][] = [];
  const fnRe = /\bfn\s+([A-Za-z_]\w*)\s*\(/g;
  let fn: RegExpExecArray | null;
  while ((fn = fnRe.exec(code))) {
    const open = code.indexOf("{", fnRe.lastIndex);
    if (open === -1) continue;
    let depth = 0;
    let i = open;
    for (; i < code.length; i++) {
      if (code[i] === "{") depth++;
      else if (code[i] === "}" && --depth === 0) break;
    }
    bodies.set(fn[1]!, code.slice(open + 1, i).match(/[A-Za-z_]\w*/g) ?? []);
    spans.push([open, i]);
    fnRe.lastIndex = i; // WGSL has no nested fns; resume past the body
  }

  // Pass 2: keep only the overrides used and functions called, dropping every
  // unrelated identifier so the retained graph stays small.
  const graph = new Map<string, FunctionNode>();
  for (const [name, identifiers] of bodies) {
    const uses = new Set<string>();
    const calls = new Set<string>();
    for (const id of identifiers) {
      if (names.has(id)) uses.add(id);
      else if (bodies.has(id)) calls.add(id);
    }
    graph.set(name, { uses, calls });
  }

  // Module scope = code with function bodies and the *declarations* themselves
  // removed, so any surviving name is a genuine module-scope use.
  let moduleScope = "";
  let cursor = 0;
  for (const [open, close] of spans) {
    moduleScope += code.slice(cursor, open);
    cursor = close + 1;
  }
  moduleScope += code.slice(cursor);
  moduleScope = moduleScope.replace(declarations, "");
  const moduleScopeUsed = new Set(
    (moduleScope.match(/[A-Za-z_]\w*/g) ?? []).filter((id) => names.has(id)),
  );

  return { names, graph, moduleScopeUsed };
}

// Names an entry point can reach: its own uses, those of every function it
// transitively calls, plus module-scope uses.
function reachable(
  { graph, moduleScopeUsed }: ReachabilityAnalysis,
  entryPoint: string,
): Set<string> {
  const used = new Set(moduleScopeUsed);
  const visited = new Set<string>();
  const stack = [entryPoint];
  while (stack.length) {
    const name = stack.pop()!;
    if (visited.has(name)) continue;
    visited.add(name);
    const node = graph.get(name);
    if (!node) continue;
    for (const name of node.uses) used.add(name);
    for (const callee of node.calls)
      if (!visited.has(callee)) stack.push(callee);
  }
  return used;
}

/**
 * Per-entry-point overrides for a (comment-stripped) source. WebKit fails Metal
 * library creation when a stage is handed a constant for an override its entry
 * point never uses — and a combined vertex+fragment source declares both
 * stages' overrides — so per-stage constants are narrowed to these sets (Chrome
 * ignores the extras). Only the small result sets are retained on the
 * reflection; the intermediate call graph is discarded.
 */
function overridesByEntryPoint(
  code: string,
  entryPoints: ShaderReflection["entryPoints"],
): Map<string, Set<string>> {
  const analysis = analyzeReachability(
    code,
    new Set(
      [...code.matchAll(/\boverride\s+([A-Za-z_]\w*)/g)].map(
        (match) => match[1]!,
      ),
    ),
    OVERRIDE_DECLARATION,
  );
  const overrides = new Map<string, Set<string>>();
  for (const entryPoint of Object.values(entryPoints)) {
    if (entryPoint) overrides.set(entryPoint, reachable(analysis, entryPoint));
  }
  return overrides;
}

const STAGE_FLAGS = {
  vertex: GPUShaderStage.VERTEX,
  fragment: GPUShaderStage.FRAGMENT,
  compute: GPUShaderStage.COMPUTE,
} as const;

/**
 * Narrow storage bindings to the stages whose entry point actually reaches
 * them.
 *
 * One source serving both stages would otherwise claim every binding in both,
 * and storage buffers are the scarce budget: `maxStorageBuffersPerShaderStage`
 * defaults to 8, and compatibility mode allows none at all in the vertex stage.
 * Restricted to storage for that reason — textures and uniforms have room to
 * spare, and a wrong answer here costs a validation error rather than a slot.
 *
 * A binding no entry point reaches keeps its original visibility: that is dead
 * code, not a stage assignment, and zero visibility is a worse thing to hand a
 * layout than a redundant flag.
 */
function scopeStorageVisibility(
  code: string,
  entryPoints: ShaderReflection["entryPoints"],
  bindings: ShaderBinding[],
): ShaderBinding[] {
  const isStorage = (binding: ShaderBinding) =>
    binding.kind === "storage" || binding.kind === "read-only-storage";

  const names = new Set(bindings.filter(isStorage).map((b) => b.name));
  if (!names.size) return bindings;

  const analysis = analyzeReachability(code, names, BINDING_DECLARATION);
  const stages = new Map<string, number>();
  for (const [stage, entryPoint] of Object.entries(entryPoints)) {
    if (!entryPoint) continue;
    const flag = STAGE_FLAGS[stage as keyof typeof STAGE_FLAGS];
    for (const name of reachable(analysis, entryPoint)) {
      stages.set(name, (stages.get(name) ?? 0) | flag);
    }
  }

  return bindings.map((binding) => {
    if (!isStorage(binding)) return binding;
    const visibility = stages.get(binding.name);
    return visibility
      ? { ...binding, visibility: visibility & binding.visibility }
      : binding;
  });
}

// ─── Public API ──────────────────────────────────────────────────────────────

/** Contents of every balanced parenthesis group, innermost first. */
function* parenGroups(source: string): Generator<string> {
  const stack: number[] = [];
  for (let i = 0; i < source.length; i++) {
    const char = source[i];
    if (char === "(") stack.push(i);
    else if (char === ")") {
      const start = stack.pop();
      if (start !== undefined) yield source.slice(start + 1, i);
    }
  }
}

// WebGPU forbids a filtering sampler on a depth (or unfilterable-float) texture.
// A plain sampler that non-comparison-samples such a texture must be declared
// "non-filtering". Sampling usually happens through helper functions (the
// textureSampleLevel call sees only parameter names), so we detect the pairing
// at call sites via argument adjacency: a texture and its sampler are always
// passed in that order (the `tex, texSampler` convention), so a plain sampler
// whose preceding argument is a depth texture binding is marked non-filtering.
// Comparison samplers are a distinct binding kind and are never marked.
//
// Argument lists are scanned depth-aware rather than by matching innermost
// parentheses: a nested call anywhere in the list — `f(tex, texSampler,
// i32(level))` — would otherwise hide the pairing, and the resulting pipeline
// fails validation with nothing pointing back at the cause.
function markNonFilteringSamplers(
  source: string,
  bindings: ShaderBinding[],
): ShaderBinding[] {
  const depthTextures = new Set(
    bindings
      .filter((b) => b.kind === "texture" && b.sampleType === "depth")
      .map((b) => b.name),
  );
  const samplers = new Set(
    bindings.filter((b) => b.kind === "sampler").map((b) => b.name),
  );
  if (!depthTextures.size || !samplers.size) return bindings;

  const nonFiltering = new Set<string>();
  for (const args of parenGroups(source)) {
    // First identifier of each argument (eg. `uPointLights[0]` → `uPointLights`).
    const idents = splitTopLevel(args, "([", ")]").map(
      (arg) => /[A-Za-z_]\w*/.exec(arg)?.[0],
    );
    for (let i = 1; i < idents.length; i++) {
      if (samplers.has(idents[i]!) && depthTextures.has(idents[i - 1]!)) {
        nonFiltering.add(idents[i]!);
      }
    }
  }
  return nonFiltering.size
    ? bindings.map((b) =>
        b.kind === "sampler" && nonFiltering.has(b.name)
          ? { ...b, nonFiltering: true }
          : b,
      )
    : bindings;
}

/** Group bindings by `@group` index, sorted by binding number within each group. */
function bindingsToGroups(
  bindings: readonly ShaderBinding[],
): Map<number, ShaderBinding[]> {
  const groups = new Map<number, ShaderBinding[]>();
  for (const binding of bindings) {
    groups.getOrInsertComputed(binding.group, () => []).push(binding);
  }
  for (const group of groups.values())
    group.sort((a, b) => a.binding - b.binding);
  return groups;
}

// Cached by source string. Bounded in practice by the number of unique
// shaders an application creates; entries are shared across contexts.
const reflectionCache = shared(
  "reflectionCache",
  () => new Map<string, ShaderReflection>(),
);

/**
 * Parse a WGSL source string into a {@link ShaderReflection}. Results are cached
 * by source identity: parsing the same string twice is free.
 */
export function parseWGSL(source: string): ShaderReflection {
  const cached = reflectionCache.get(source);
  if (cached) return cached;

  const stripped = stripComments(source);
  const rawStructs = parseStructs(stripped);
  const structs = new Map<string, StructLayout>();
  const { entryPoints, vertexParams } = parseEntryPoints(stripped);

  let visibility = 0;
  if (entryPoints.vertex) visibility |= GPUShaderStage.VERTEX;
  if (entryPoints.fragment) visibility |= GPUShaderStage.FRAGMENT;
  if (entryPoints.compute) visibility |= GPUShaderStage.COMPUTE;

  const bindings = scopeStorageVisibility(
    stripped,
    entryPoints,
    markNonFilteringSamplers(
      stripped,
      parseBindings(stripped, structs, rawStructs, visibility),
    ),
  );
  const vertexInputs = vertexParams.length
    ? parseVertexInputs(vertexParams[0]!, rawStructs)
    : [];

  const reflection: ShaderReflection = {
    bindings,
    groups: bindingsToGroups(bindings),
    vertexInputs,
    structs,
    entryPoints,
    overrides: overridesByEntryPoint(stripped, entryPoints),
  };
  reflectionCache.set(source, reflection);
  return reflection;
}

/**
 * {@link parseWGSL} accounted against a device's `shaderReflections` cache
 * counters — what the package itself calls, so the cost of reflecting a new
 * shader is visible next to compiling one. `parseWGSL` stays device-free:
 * reflection is pure and runs outside a WebGPU host.
 */
export function reflectWGSL(
  device: GPUDevice,
  source: string,
): ShaderReflection {
  const cached = reflectionCache.get(source);
  if (cached) {
    recordCacheAccess(device, "shaderReflections", null);
    return cached;
  }
  debugLog(device, "resources", () => `parseWGSL (${source.length} chars)`);
  const start = performance.now();
  const reflection = parseWGSL(source);
  recordCacheAccess(device, "shaderReflections", performance.now() - start);
  return reflection;
}

/**
 * The reflection for one pipeline's stages: {@link reflectWGSL} per source plus
 * {@link mergeReflections} when they differ, charged to the device's
 * `shaderReflections` counters as a single access. The merge is not cached and
 * runs per pipeline, so it belongs in the same number as the parses it joins.
 */
export function reflectStages(
  device: GPUDevice,
  vertex: string,
  fragment?: string,
): ShaderReflection {
  if (!fragment || fragment === vertex) return reflectWGSL(device, vertex);
  const cachedVertex = reflectionCache.get(vertex);
  const cachedFragment = reflectionCache.get(fragment);
  if (cachedVertex && cachedFragment) {
    recordCacheAccess(device, "shaderReflections", null);
    return mergeReflections(cachedVertex, cachedFragment);
  }
  debugLog(
    device,
    "resources",
    () =>
      `parseWGSL vertex+fragment (${vertex.length}+${fragment.length} chars)`,
  );
  const start = performance.now();
  const reflection = mergeReflections(
    cachedVertex ?? parseWGSL(vertex),
    cachedFragment ?? parseWGSL(fragment),
  );
  recordCacheAccess(device, "shaderReflections", performance.now() - start);
  return reflection;
}

/**
 * Merge reflections of separate vertex/fragment sources: bindings sharing a
 * group/binding slot are unified with combined stage visibility.
 */
export function mergeReflections(
  a: ShaderReflection,
  b: ShaderReflection,
): ShaderReflection {
  if (a === b) return a;

  const bindings: ShaderBinding[] = [...a.bindings];
  for (const binding of b.bindings) {
    const index = bindings.findIndex(
      (existing) =>
        existing.group === binding.group &&
        existing.binding === binding.binding,
    );
    if (index === -1) {
      bindings.push(binding);
    } else {
      const existing = bindings[index]!;
      if (existing.kind !== binding.kind) {
        throw reflectionError(
          `binding "${binding.name}" (@group(${binding.group}) @binding(${binding.binding})) has conflicting types across shader stages`,
        );
      }
      bindings[index] = {
        ...existing,
        visibility: existing.visibility | binding.visibility,
      };
    }
  }

  return {
    bindings,
    groups: bindingsToGroups(bindings),
    vertexInputs: a.vertexInputs.length ? a.vertexInputs : b.vertexInputs,
    structs: new Map([...a.structs, ...b.structs]),
    entryPoints: { ...b.entryPoints, ...a.entryPoints },
    // Entry point names are distinct per stage, so the maps never collide.
    overrides: new Map([...a.overrides, ...b.overrides]),
  };
}