import { packStruct } from "./pack.js";
import { isGpuBuffer } from "./buffer.js";
import { isGpuTexture } from "./texture.js";
import { debugCounters, debugLog, recordCacheAccess } from "./debug.js";
import {
  commandsState,
  trackBindGroupResource,
  RingAllocator,
} from "./state.js";
import { objectId } from "./shared.js";
import type { PackableValue } from "./pack.js";
import type {
  GpuContext,
  ShaderBinding,
  ShaderReflection,
  Uniforms,
} from "./types.js";

const missingUniform = (binding: ShaderBinding, detail = "") =>
  new Error(
    `pex-gpu: missing uniform "${binding.name}" (@group(${binding.group}) @binding(${binding.binding}))${detail}`,
  );

// ─── Bind group layouts ──────────────────────────────────────────────────────

function layoutEntry(binding: ShaderBinding): GPUBindGroupLayoutEntry {
  const base = { binding: binding.binding, visibility: binding.visibility };
  switch (binding.kind) {
    case "uniform":
      return {
        ...base,
        buffer: {
          type: "uniform",
          hasDynamicOffset: true,
          ...(binding.minBindingSize && {
            minBindingSize: binding.minBindingSize,
          }),
        },
      };
    case "storage":
    case "read-only-storage":
      return {
        ...base,
        buffer: {
          type: binding.kind,
          // A runtime array filled from a JS array is packed into the shared
          // ring, so its slice is addressed by dynamic offset — same as a
          // uniform. A caller-supplied GpuBuffer binds at offset 0.
          ...(binding.arrayElementLayout && {
            hasDynamicOffset: true,
            minBindingSize: binding.minBindingSize,
          }),
        },
      };
    case "sampler":
      return {
        ...base,
        sampler: { type: binding.nonFiltering ? "non-filtering" : "filtering" },
      };
    case "comparison-sampler":
      return { ...base, sampler: { type: "comparison" } };
    case "texture":
      return {
        ...base,
        texture: {
          sampleType: binding.sampleType ?? "float",
          viewDimension: binding.viewDimension ?? "2d",
          multisampled: binding.multisampled ?? false,
        },
      };
    case "storage-texture":
      return {
        ...base,
        storageTexture: {
          // reflection.ts always sets storageTextureFormat for this kind.
          format: binding.storageTextureFormat!,
          access: binding.storageTextureAccess ?? "write-only",
          viewDimension: binding.viewDimension ?? "2d",
        },
      };
    case "external-texture":
      return { ...base, externalTexture: {} };
  }
}

const layoutKeyFor = (bindings: readonly ShaderBinding[]): string =>
  bindings
    .map(
      (binding) =>
        `${binding.binding}:${binding.kind}:${binding.visibility}:${binding.sampleType ?? ""}:${binding.viewDimension ?? ""}:${binding.multisampled ?? ""}:${binding.storageTextureFormat ?? ""}:${binding.minBindingSize ?? ""}:${binding.nonFiltering ?? ""}`,
    )
    .join("|");

function bindGroupLayout(
  ctx: GpuContext,
  bindings: readonly ShaderBinding[],
): { layout: GPUBindGroupLayout; key: string } {
  const state = commandsState(ctx);
  const key = layoutKeyFor(bindings);
  let created = false;
  let start = 0;

  const layout = state.bindGroupLayouts.getOrInsertComputed(key, () => {
    created = true;
    start = performance.now();
    debugLog(ctx.device, "resources", () => `createBindGroupLayout [${key}]`);
    return ctx.device.createBindGroupLayout({
      label: "pex-gpu bind group layout",
      entries: bindings.map(layoutEntry),
    });
  });
  recordCacheAccess(
    ctx.device,
    "bindGroupLayouts",
    created ? performance.now() - start : null,
  );
  return { layout, key };
}

/**
 * Explicit pipeline layout from reflected bindings. Shared across pipelines
 * with matching group layouts, so bind groups are reusable between them (unlike
 * `layout: "auto"`).
 */
export function pipelineLayout(
  ctx: GpuContext,
  reflection: ShaderReflection,
): { layout: GPUPipelineLayout; key: string } {
  const state = commandsState(ctx);
  // Group indices need not be contiguous or start at 0 — iterate up to the
  // highest index used (0 groups if there are none).
  const groupCount = Math.max(-1, ...reflection.groups.keys()) + 1;

  const keys: string[] = [];
  const layouts: GPUBindGroupLayout[] = [];
  for (let index = 0; index < groupCount; index++) {
    const bindings = reflection.groups.get(index) ?? [];
    const { layout, key } = bindGroupLayout(ctx, bindings);
    keys.push(key);
    layouts.push(layout);
  }

  const key = keys.join("||") || "empty";
  let created = false;
  let start = 0;

  const layout = state.pipelineLayouts.getOrInsertComputed(key, () => {
    created = true;
    start = performance.now();
    debugLog(ctx.device, "resources", () => `createPipelineLayout [${key}]`);
    return ctx.device.createPipelineLayout({
      label: "pex-gpu pipeline layout",
      bindGroupLayouts: layouts,
    });
  });
  recordCacheAccess(
    ctx.device,
    "pipelineLayouts",
    created ? performance.now() - start : null,
  );
  return { layout, key };
}

// ─── Uniform value resolution ────────────────────────────────────────────────

function uniformStructValues(
  binding: ShaderBinding,
  uniforms: Uniforms,
): Record<string, PackableValue> | PackableValue {
  const direct = uniforms[binding.name];
  if (direct !== undefined) {
    return direct as Record<string, PackableValue> | PackableValue;
  }

  // Fall back to flat member names: uniforms: { projection, view, model }.
  const layout = binding.structLayout!;
  const values: Record<string, PackableValue> = {};
  const missing: string[] = [];
  for (const member of layout.members) {
    const value = uniforms[member.name];
    if (value === undefined) missing.push(member.name);
    else values[member.name] = value as PackableValue;
  }
  if (missing.length === layout.members.length) throw missingUniform(binding);
  if (missing.length) {
    throw missingUniform(binding, `: missing members ${missing.join(", ")}`);
  }
  return values;
}

// ─── Bind group building ─────────────────────────────────────────────────────

export interface ResolvedBindGroup {
  index: number;
  bindGroup: GPUBindGroup;
  dynamicOffsets: number[];
}

/**
 * Build (or reuse cached) bind groups for every reflected group. Uniform buffer
 * bindings are packed into the per-frame ring allocator and bound with dynamic
 * offsets, so bind groups are stable across draws and frames.
 */
export function buildBindGroups(
  ctx: GpuContext,
  reflection: ShaderReflection,
  uniforms: Uniforms,
  overrides?: Record<number, GPUBindGroup>,
): ResolvedBindGroup[] {
  const state = commandsState(ctx);
  // Group indices need not be contiguous or start at 0 — iterate up to the
  // highest index used (0 groups if there are none).
  const groupCount = Math.max(-1, ...reflection.groups.keys()) + 1;
  const result: ResolvedBindGroup[] = [];

  for (let index = 0; index < groupCount; index++) {
    const override = overrides?.[index];
    if (override) {
      result.push({ index, bindGroup: override, dynamicOffsets: [] });
      continue;
    }

    const bindings = reflection.groups.get(index) ?? [];
    const { layout, key: layoutKey } = bindGroupLayout(ctx, bindings);
    const entries: GPUBindGroupEntry[] = [];
    const dynamicOffsets: number[] = [];
    const keyParts: (string | number)[] = [layoutKey];
    const trackedResources: { id: number; wrapper: object }[] = [];
    let cacheable = true;

    for (const binding of bindings) {
      const value = uniforms[binding.name];

      switch (binding.kind) {
        case "uniform": {
          // reflection.ts always sets structLayout for uniform bindings.
          const layout = binding.structLayout!;
          const values = uniformStructValues(binding, uniforms);
          const allocator = (state.allocator ??= new RingAllocator(ctx.device));
          const offset = allocator.allocate(layout.size);
          debugCounters(ctx.device).frame.uniformBytes += layout.size;
          packStruct(allocator.staging, layout, values, offset);
          entries.push({
            binding: binding.binding,
            resource: { buffer: allocator.buffer, size: layout.size },
          });
          dynamicOffsets.push(offset);
          keyParts.push(objectId(allocator.buffer));
          break;
        }
        case "storage":
        case "read-only-storage": {
          if (binding.arrayElementLayout && !isGpuBuffer(value)) {
            if (!Array.isArray(value)) {
              throw missingUniform(
                binding,
                " — expected an array or a storage GpuBuffer",
              );
            }
            const element = binding.arrayElementLayout;
            const stride = binding.arrayStride!;
            // A binding cannot be zero-sized, so an empty array still reserves
            // one element; the shader reads none of it, having the count.
            const count = Math.max(value.length, 1);
            const size = stride * count;
            const allocator = (state.storageAllocator ??= new RingAllocator(
              ctx.device,
              "read-only-storage",
            ));
            const offset = allocator.allocate(size);
            debugCounters(ctx.device).frame.uniformBytes += size;
            for (let i = 0; i < value.length; i++) {
              packStruct(
                allocator.staging,
                element,
                value[i] as unknown as Record<string, PackableValue>,
                offset + i * stride,
              );
            }
            entries.push({
              binding: binding.binding,
              resource: { buffer: allocator.buffer, size },
            });
            dynamicOffsets.push(offset);
            // Size varies with the count, so it belongs in the key: the same
            // buffer bound at a different length is a different bind group.
            keyParts.push(objectId(allocator.buffer), size);
            break;
          }
          if (!isGpuBuffer(value)) {
            throw missingUniform(binding, " — expected a storage GpuBuffer");
          }
          if (binding.arrayElementLayout) dynamicOffsets.push(0);
          entries.push({
            binding: binding.binding,
            resource: { buffer: value.buffer },
          });
          keyParts.push(value.id);
          trackedResources.push({ id: value.id, wrapper: value });
          break;
        }
        case "sampler":
        case "comparison-sampler": {
          if (!(value instanceof GPUSampler)) {
            throw missingUniform(binding, " — expected a GPUSampler");
          }
          entries.push({ binding: binding.binding, resource: value });
          keyParts.push(objectId(value));
          break;
        }
        case "texture":
        case "storage-texture": {
          const view = isGpuTexture(value)
            ? value.view
            : value instanceof GPUTextureView
              ? value
              : null;
          if (!view) {
            throw missingUniform(
              binding,
              " — expected a GpuTexture or GPUTextureView",
            );
          }
          entries.push({ binding: binding.binding, resource: view });
          keyParts.push(objectId(view));
          if (isGpuTexture(value)) {
            trackedResources.push({ id: value.id, wrapper: value });
          }
          break;
        }
        case "external-texture": {
          if (!(value instanceof HTMLVideoElement)) {
            throw missingUniform(binding, " — expected an HTMLVideoElement");
          }
          entries.push({
            binding: binding.binding,
            resource: ctx.device.importExternalTexture({ source: value }),
          });
          cacheable = false; // external textures expire every frame
          break;
        }
      }
    }

    let bindGroup: GPUBindGroup | undefined;
    const cacheKey = keyParts.join(",");
    if (cacheable) bindGroup = state.bindGroups.get(cacheKey);
    if (bindGroup) {
      recordCacheAccess(ctx.device, "bindGroups", null);
    } else {
      // External-texture groups rebuild every frame by design: not worth logging
      if (cacheable) {
        debugLog(
          ctx.device,
          "resources",
          () => `createBindGroup [${cacheKey}]`,
        );
      }
      const start = performance.now();
      bindGroup = ctx.device.createBindGroup({
        label: "pex-gpu bind group",
        layout,
        entries,
      });
      recordCacheAccess(ctx.device, "bindGroups", performance.now() - start);
      if (cacheable) state.bindGroups.set(cacheKey, bindGroup);
    }
    if (cacheable) {
      for (const { id, wrapper } of trackedResources) {
        trackBindGroupResource(state, id, wrapper, cacheKey);
      }
    }
    result.push({ index, bindGroup, dynamicOffsets });
  }

  return result;
}