import { BUFFER_USAGE_PRESETS } from "./buffer.js";
import { debugLog, trackResource, untrackResource } from "./debug.js";
import { assertSchema, objectId, shared } from "./shared.js";
import type {
  ComputePipeline,
  GpuContext,
  GpuTexture,
  RenderPipeline,
  ShaderReflection,
} from "./types.js";

/** How a {@link RingAllocator} is bound: sets its usage and offset alignment. */
export type RingAllocatorKind = "uniform" | "read-only-storage";

/**
 * Per-frame ring buffer. Each draw's structs are packed into CPU staging at
 * aligned offsets, bound through bind groups with dynamic offsets, and flushed
 * with a single writeBuffer before submission.
 *
 * One instance per binding kind: uniform and storage have different usage flags
 * and different minimum offset alignments, and a slice has to satisfy the one
 * it will be bound as.
 */
export class RingAllocator {
  buffer: GPUBuffer;
  /** Incremented when the GPU buffer is replaced; part of bind group cache keys. */
  generation = 0;
  #device: GPUDevice;
  #kind: RingAllocatorKind;
  #staging: ArrayBuffer;
  #cursor = 0;
  #alignment: number;
  #retired: GPUBuffer[] = [];

  constructor(
    device: GPUDevice,
    kind: RingAllocatorKind = "uniform",
    initialSize = 256 * 1024,
  ) {
    this.#device = device;
    this.#kind = kind;
    this.#alignment =
      kind === "uniform"
        ? device.limits.minUniformBufferOffsetAlignment
        : device.limits.minStorageBufferOffsetAlignment;
    this.#staging = new ArrayBuffer(initialSize);
    this.buffer = this.#createBuffer(initialSize);
  }

  #createBuffer(size: number): GPUBuffer {
    const buffer = this.#device.createBuffer({
      label: `pex-gpu ${this.#kind === "uniform" ? "uniforms" : "storage"}`,
      size,
      usage: BUFFER_USAGE_PRESETS[this.#kind],
    });
    trackResource(this.#device, "buffers", {
      id: objectId(buffer),
      label: buffer.label,
      usage: buffer.usage,
      bytes: size,
    });
    return buffer;
  }

  #destroyBuffer(buffer: GPUBuffer): void {
    untrackResource(this.#device, "buffers", objectId(buffer));
    buffer.destroy();
  }

  get staging(): ArrayBuffer {
    return this.#staging;
  }

  /** Reserve an aligned slice; returns its byte offset. Grows when exhausted. */
  allocate(size: number): number {
    const offset = Math.ceil(this.#cursor / this.#alignment) * this.#alignment;
    if (offset + size > this.#staging.byteLength) {
      this.#grow(Math.max(this.#staging.byteLength * 2, offset + size));
      return this.allocate(size);
    }
    this.#cursor = offset + size;
    // A member the caller left out reads as zero rather than as whatever the
    // last draw staged at this offset. Struct shapes that are only partly
    // written — a union packed from one of several light types, say — depend
    // on it, and a stale default is otherwise invisible until it isn't.
    new Uint8Array(this.#staging, offset, size).fill(0);
    return offset;
  }

  #grow(newSize: number) {
    // Flush what's staged so far into the current buffer — bind groups
    // recorded against it stay valid for this frame — then start a new one.
    this.flush();
    this.#retired.push(this.buffer);
    this.buffer = this.#createBuffer(newSize);
    this.generation++;
    debugLog(
      this.#device,
      "resources",
      () =>
        `${this.#kind} ring grown to ${newSize} bytes (generation ${this.generation})`,
    );
    this.#staging = new ArrayBuffer(newSize);
    this.#cursor = 0;
  }

  /** Upload the staged range. Called before queue submission. */
  flush(): void {
    if (this.#cursor === 0) return;
    this.#device.queue.writeBuffer(
      this.buffer,
      0,
      this.#staging,
      0,
      this.#cursor,
    );
  }

  /** Start a new frame: reclaim the cursor and destroy retired buffers. */
  reset(): void {
    this.#cursor = 0;
    for (const buffer of this.#retired) this.#destroyBuffer(buffer);
    this.#retired.length = 0;
  }

  dispose(): void {
    this.reset();
    this.#destroyBuffer(this.buffer);
  }
}

export interface ActivePass {
  encoder: GPURenderPassEncoder;
  colorFormats: (GPUTextureFormat | null)[];
  depthStencilFormat: GPUTextureFormat | undefined;
  /** Pipelines in this pass must not write depth. */
  depthReadOnly: boolean;
  sampleCount: number;
  /** Signature of everything above, part of pipeline variant cache keys. */
  key: string;
}

export interface ActiveComputePass {
  encoder: GPUComputePassEncoder;
}

/**
 * Per-pipeline-object cache: shader reflection and pipeline layout are resolved
 * once per plain pipeline object (its WGSL sources are assumed immutable), and
 * GPU pipelines are stored per variant signature — pass formats, vertex layout
 * and mutable pipeline state.
 */
export interface CachedPipelines<T> {
  reflection: ShaderReflection;
  layout: GPUPipelineLayout;
  layoutKey: string;
  variants: Map<string, T>;
}

/**
 * Adds the per-stage sources resolved from the pipeline's shader form, so the
 * combined `shader` is unfolded once rather than at every variant miss.
 * `fragmentSource` undefined means a pipeline with no fragment stage.
 */
export interface CachedRenderPipelines extends CachedPipelines<GPURenderPipeline> {
  vertexSource: string;
  fragmentSource: string | undefined;
}

export interface FrameState {
  encoder: GPUCommandEncoder;
  /**
   * The literal swapchain texture for this segment, distinct from
   * FrameSegment.canvasTexture when preserveDrawingBuffer routes rendering
   * through a backbuffer instead — endFrame() blits into this.
   */
  swapchainTexture: GPUTexture;
  canvasView: GPUTextureView;
  /**
   * Whether the managed canvas depth/stencil has been attached to a pass this
   * frame — it clears at first use and loads afterwards.
   */
  canvasDepthStencilCleared: boolean;
  activePass: ActivePass | null;
  activeComputePass: ActiveComputePass | null;
}

export interface CommandsState {
  /**
   * Lazy — created at the first submit that packs uniforms, so the frame loop
   * costs nothing for code that never submits a declarative command.
   */
  allocator: RingAllocator | null;
  /** Runtime-sized `array<T>` storage bindings filled from a JS array. */
  storageAllocator: RingAllocator | null;
  bindGroupLayouts: Map<string, GPUBindGroupLayout>;
  pipelineLayouts: Map<string, GPUPipelineLayout>;
  /**
   * Cache keys embed the ids of the GpuBuffer/GpuTexture wrappers bound in each
   * group; a disposed-and-replaced resource would otherwise leave its old bind
   * group cached forever. `bindGroupsByResource` is the reverse index (resource
   * id → cache keys), pruned by `bindGroupFinalizer` once a wrapper is
   * unreachable — GC-triggered rather than tied to `dispose()`, so a
   * GpuBuffer/GpuTexture's dispose() never needs to know the
   * declarative-commands cache exists.
   */
  bindGroups: Map<string, GPUBindGroup>;
  bindGroupsByResource: Map<number, Set<string>>;
  bindGroupFinalizer: FinalizationRegistry<number>;
  renderPipelines: WeakMap<RenderPipeline, CachedRenderPipelines>;
  computePipelines: WeakMap<
    ComputePipeline,
    CachedPipelines<GPUComputePipeline>
  >;
  /**
   * Pipeline objects whose reflection threw or whose GPU pipeline failed
   * validation. Their submits are skipped: the error is reported once, and an
   * invalid pipeline would otherwise invalidate the whole command buffer.
   */
  failedPipelines: WeakSet<RenderPipeline | ComputePipeline>;
  /**
   * Context-managed canvas-sized depth/stencil texture (depthStencilAttachment:
   * "canvas").
   */
  canvasDepthStencil: GpuTexture | null;
  frame: FrameState | null;
}

const states = shared(
  "commandsState",
  () => new WeakMap<GpuContext, CommandsState>(),
);

export function commandsState(ctx: GpuContext): CommandsState {
  return states.getOrInsertComputed(ctx, () => {
    // Only on the miss path: a context already in the map is this copy's by
    // construction, so steady-state submits pay nothing for the check.
    assertSchema(ctx);
    const bindGroups = new Map<string, GPUBindGroup>();
    const bindGroupsByResource = new Map<number, Set<string>>();
    const bindGroupFinalizer = new FinalizationRegistry<number>((id) => {
      const keys = bindGroupsByResource.get(id);
      if (!keys) return;
      for (const key of keys) bindGroups.delete(key);
      bindGroupsByResource.delete(id);
    });
    return {
      allocator: null,
      storageAllocator: null,
      bindGroupLayouts: new Map(),
      pipelineLayouts: new Map(),
      bindGroups,
      bindGroupsByResource,
      bindGroupFinalizer,
      renderPipelines: new WeakMap(),
      computePipelines: new WeakMap(),
      failedPipelines: new WeakSet(),
      canvasDepthStencil: null,
      frame: null,
    };
  });
}

/**
 * Like {@link commandsState}, but never creates one — for teardown paths that
 * must not allocate state for a context that never submitted a declarative
 * command.
 */
export function peekCommandsState(ctx: GpuContext): CommandsState | undefined {
  return states.get(ctx);
}

/**
 * Register a bind-group-cacheable resource wrapper (a GpuBuffer/GpuTexture, not
 * the raw GPUBuffer/GPUTextureView the bind group actually retains — a
 * FinalizationRegistry watching the retained object would never fire, since the
 * very cache entry it's meant to prune keeps it alive) so its cache entries are
 * pruned once the wrapper is unreachable. A no-op past the first call for a
 * given id.
 */
export function trackBindGroupResource(
  state: CommandsState,
  id: number,
  wrapper: object,
  cacheKey: string,
): void {
  let keys = state.bindGroupsByResource.get(id);
  if (!keys) {
    keys = new Set();
    state.bindGroupsByResource.set(id, keys);
    state.bindGroupFinalizer.register(wrapper, id);
  }
  keys.add(cacheKey);
}

export function frameState(ctx: GpuContext): FrameState {
  const state = commandsState(ctx).frame;
  if (!state) {
    throw new Error(
      "pex-gpu: submit() must be called between beginFrame() and endFrame() eg. frame() callback",
    );
  }
  return state;
}