import type { PipelineCache } from "./pipeline-cache.js";
import type { PackableValue } from "./pack.js";

export type TypedArray =
  | Int8Array
  | Uint8Array
  | Uint8ClampedArray
  | Int16Array
  | Uint16Array
  | Int32Array
  | Uint32Array
  | Float16Array
  | Float32Array
  | Float64Array
  | BigInt64Array
  | BigUint64Array;

/** Data accepted by buffer and texture factories. Nested arrays are flattened. */
export type BufferData = TypedArray | ArrayBuffer | number[] | number[][];

// ─── Context ─────────────────────────────────────────────────────────────────

/**
 * Aggregates the adapter/device/canvas-context triple WebGPU leaves separate,
 * plus per-context pex-gpu state.
 */
export interface GpuContext extends Disposable {
  readonly adapter: GPUAdapter;
  readonly device: GPUDevice;
  readonly canvas: HTMLCanvasElement;
  readonly canvasContext: GPUCanvasContext;
  /**
   * Preferred canvas format, used as the color target format when rendering to
   * the canvas.
   */
  readonly format: GPUTextureFormat;
  /**
   * How the compositor reads the canvas' alpha, as configured. Retained because
   * `GPUCanvasContext` exposes no getter for it, and anything drawing the final
   * image has to agree: `"premultiplied"` reads the stored values as already
   * multiplied by alpha, so a renderer writing straight alpha composites too
   * bright, while `"opaque"` ignores alpha entirely.
   */
  readonly alphaMode: GPUCanvasAlphaMode;
  /**
   * Format of the context-managed canvas depth/stencil texture, resolved from
   * the `depth`/`stencil`/`depthStencilFormat` context options — undefined when
   * disabled. Use as the depthStencil format when building pipelines or render
   * bundles targeting the canvas by hand.
   */
  readonly depthStencilFormat: GPUTextureFormat | undefined;
  readonly pipelineCache: PipelineCache;
  /**
   * Physical pixels per CSS pixel, as applied by {@link resize}: the factor
   * mapping CSS coordinates to drawing buffer ones.
   */
  pixelRatio: number;
  /** Drawing buffer width in physical pixels. */
  width: number;
  /** Drawing buffer height in physical pixels. */
  height: number;
  /**
   * Release the device/canvas and declarative-commands state. Idempotent;
   * identical to `[Symbol.dispose]`, for call sites without `using`.
   */
  dispose(): void;
}

/**
 * A recording segment opened by {@link beginFrame} and closed by
 * {@link endFrame}: the scope submit() and other declarative commands record
 * into. `frame()`'s callback receives the richer {@link FrameInfo}, which adds
 * requestAnimationFrame-loop bookkeeping on top of this.
 */
export interface FrameSegment {
  /**
   * This segment's command encoder. Finished and submitted automatically by
   * endFrame().
   */
  readonly encoder: GPUCommandEncoder;
  /**
   * The canvas render target for this segment. With `preserveDrawingBuffer`
   * enabled, this is the persistent backing texture, not the literal swapchain
   * texture — it's copied into the canvas automatically by endFrame(), so its
   * content (unlike the raw swapchain texture's) survives into the next
   * segment.
   */
  readonly canvasTexture: GPUTexture;
  readonly canvasView: GPUTextureView;
  readonly width: number;
  readonly height: number;
}

export interface FrameInfo extends FrameSegment {
  /** Seconds since the frame loop started. */
  readonly time: number;
  /** Seconds since the previous frame. */
  readonly deltaTime: number;
  /** Frame index since the frame loop started; 0 on the first frame. */
  readonly frame: number;
  /**
   * Raw requestAnimationFrame timestamp: milliseconds on the performance.now()
   * timeline, for correlating frames with event.timeStamp, VideoFrame callbacks
   * or performance marks.
   */
  readonly timestamp: DOMHighResTimeStamp;
  /** True on frames following a resize() call. */
  readonly resized: boolean;
}

// ─── Resources ───────────────────────────────────────────────────────────────

/** Named usage presets resolving to common GPUBufferUsage flag combinations. */
export type BufferUsagePreset =
  | "vertex"
  | "index"
  | "indirect"
  | "uniform"
  | "storage"
  | "read-only-storage"
  | "upload"
  | "readback"
  | "query-resolve";

/** Buffer resource wrapping a GPUBuffer. */
export interface GpuBuffer
  extends Disposable, Pick<GPUBuffer, "size" | "usage"> {
  readonly buffer: GPUBuffer;
  /**
   * Element count of the source data, when created from an array. Kept in sync
   * by updateBuffer, so a draw inferring its vertex or index count from it sees
   * the contents rather than the allocation.
   */
  length?: number;
  /** Index format, when usage includes GPUBufferUsage.INDEX. */
  readonly indexFormat?: GPUIndexFormat;
  /** Unique id used for bind group caching. */
  readonly id: number;
  /**
   * Release the GPU buffer. Idempotent; identical to `[Symbol.dispose]`, for
   * call sites without `using`.
   */
  dispose(): void;
}

/** Texture resource wrapping a GPUTexture. */
export interface GpuTexture
  extends
    Disposable,
    Pick<
      GPUTexture,
      | "format"
      | "width"
      | "height"
      | "depthOrArrayLayers"
      | "mipLevelCount"
      | "sampleCount"
      | "dimension"
    > {
  readonly texture: GPUTexture;
  /**
   * Cached default view (dimension inferred, overridable via viewDimension
   * option).
   */
  readonly view: GPUTextureView;
  /**
   * Dimension of `view`: "cube" for a cubemap, "2d-array" for any other layered
   * texture. Distinct from `dimension`, which is the texture's own and is "2d"
   * for all three — this is what tells them apart.
   */
  readonly viewDimension: GPUTextureViewDimension;
  /** Unique id used for bind group caching. */
  readonly id: number;
  /**
   * Release the GPU texture. Idempotent; identical to `[Symbol.dispose]`, for
   * call sites without `using`.
   */
  dispose(): void;
}

// ─── WGSL reflection ─────────────────────────────────────────────────────────

export interface StructMember {
  readonly name: string;
  /**
   * Normalized WGSL type, eg. "f32", "vec3<f32>", "mat4x4<f32>",
   * "array<vec4<f32>, 8>".
   */
  readonly type: string;
  readonly offset: number;
  readonly size: number;
  readonly align: number;
  /** Present for fixed-size arrays. */
  readonly arrayCount?: number;
  /** Present for fixed-size arrays. */
  readonly arrayStride?: number;
  /**
   * Element struct layout for `array<SomeStruct, N>` members/bindings, so each
   * element is packed by member name (with its internal padding) at
   * arrayStride, rather than from tightly packed source data.
   */
  readonly elementLayout?: StructLayout;
  /**
   * Layout of a member that is itself a struct, packed from a nested object of
   * its own member names.
   *
   * Its offset is computed at a 16-byte alignment, the std140-style reading of
   * the uniform address space rules — but implementations disagree (Dawn uses
   * the struct's natural alignment), and nothing here knows the address space
   * to pick between them. So a nested struct member whose own alignment is
   * under 16 must carry an explicit `@align` in the WGSL, or what is packed and
   * what is read come apart silently.
   */
  readonly structLayout?: StructLayout;
}

export interface StructLayout {
  readonly name: string;
  /** Total size including trailing padding. */
  readonly size: number;
  readonly align: number;
  readonly members: readonly StructMember[];
}

export type BindingKind =
  | "uniform"
  | "storage"
  | "read-only-storage"
  | "sampler"
  | "comparison-sampler"
  | "texture"
  | "storage-texture"
  | "external-texture";

export interface ShaderBinding {
  readonly name: string;
  readonly group: number;
  readonly binding: number;
  readonly kind: BindingKind;
  readonly visibility: GPUShaderStageFlags;
  /** Struct layout for uniform/storage buffer bindings on struct types. */
  readonly structLayout?: StructLayout;
  /** Byte size for buffer bindings with a known fixed size. */
  readonly minBindingSize?: number;
  /**
   * Element layout and stride of a runtime-sized `array<T>` storage binding, so
   * its value can be supplied as a JS array and packed like a uniform
   * array-of-struct is. Absent for fixed-size arrays and non-array storage,
   * whose size is the caller's to own — those take a GpuBuffer.
   */
  readonly arrayElementLayout?: StructLayout;
  readonly arrayStride?: number;
  /** For texture bindings. */
  readonly sampleType?: GPUTextureSampleType;
  /** For texture and storage-texture bindings. */
  readonly viewDimension?: GPUTextureViewDimension;
  /** For texture bindings. */
  readonly multisampled?: boolean;
  /** For storage-texture bindings. */
  readonly storageTextureFormat?: GPUTextureFormat;
  /** For storage-texture bindings. */
  readonly storageTextureAccess?: GPUStorageTextureAccess;
  /**
   * For sampler bindings used to non-comparison-sample a depth (or otherwise
   * unfilterable) texture: the layout entry must be "non-filtering".
   */
  readonly nonFiltering?: boolean;
}

export interface VertexInput {
  readonly name: string;
  readonly location: number;
  readonly type: string;
  readonly format: GPUVertexFormat;
  /** Tightly packed byte size of the attribute. */
  readonly arrayStride: number;
}

export interface ShaderReflection {
  readonly bindings: readonly ShaderBinding[];
  /** Bindings grouped by `@group` index. */
  readonly groups: ReadonlyMap<number, readonly ShaderBinding[]>;
  readonly vertexInputs: readonly VertexInput[];
  /**
   * Struct layouts keyed `"<address space>:<name>"` — a struct bound as both a
   * uniform and a storage buffer has two, since only the uniform space rounds
   * array strides and nested struct alignments up to 16.
   */
  readonly structs: ReadonlyMap<string, StructLayout>;
  readonly entryPoints: {
    vertex?: string;
    fragment?: string;
    compute?: string;
  };
  /**
   * Pipeline-overridable constants each entry point can reference — directly,
   * through the functions it transitively calls, or at module scope — keyed by
   * entry point name. Used to narrow per-stage `constants`, since WebKit
   * rejects a value for an override the stage's entry point never uses.
   */
  readonly overrides: ReadonlyMap<string, ReadonlySet<string>>;
}

// ─── Pass ────────────────────────────────────────────────────────────────────

export interface ColorAttachment {
  /**
   * Render target. Omit (or omit `colorAttachments` entirely) to render to the
   * canvas.
   */
  texture?: GpuTexture;
  /**
   * Explicit view override, eg. a cubemap face or mip level. Defaults to
   * texture.view, or the canvas.
   */
  view?: GPUTextureView;
  /** Multisample resolve target. "canvas" resolves to the swapchain. */
  resolveTarget?: GpuTexture | "canvas";
  /** Setting a clear value implies loadOp "clear". */
  clearValue?: GPUColor;
  loadOp?: GPULoadOp;
  storeOp?: GPUStoreOp;
}

/**
 * Depth/stencil attachment. `view` defaults to texture.view; setting a clear
 * value implies the matching load op ("load" otherwise); store ops default to
 * "store".
 */
export interface DepthStencilAttachment extends Partial<GPURenderPassDepthStencilAttachment> {
  texture: GpuTexture;
}

export interface RenderPassDescriptor extends GPUObjectDescriptorBase {
  /** Omit (along with `clearValue`) to render to the canvas. */
  colorAttachments?: ColorAttachment[];
  /**
   * Depth/stencil attachment. Canvas-targeting passes default to "canvas": the
   * context-managed, canvas-sized depth/stencil texture (format
   * `ctx.depthStencilFormat`, resizes automatically), so most passes never
   * create a depth/stencil texture explicitly. It clears at its first use each
   * frame (depth 1, stencil 0) and loads in later passes.
   */
  depthStencilAttachment?: GpuTexture | DepthStencilAttachment | "canvas";
  /** Shorthand for colorAttachments[0].clearValue targeting the canvas. */
  clearValue?: GPUColor;
  /** Shorthand for depthStencilAttachment.depthClearValue. */
  depthClearValue?: number;
  /** Shorthand for depthStencilAttachment.stencilClearValue. */
  stencilClearValue?: GPUStencilValue;
  /**
   * Time the pass with a {@link TimestampQuery}: `timestampWrites:
   * query.timestampWrites`.
   */
  timestampWrites?: GPURenderPassTimestampWrites;
}

/**
 * Compute pass options. Compute passes have no attachments — only a label and
 * optional timestamp writes (`timestampWrites: query.timestampWrites` with a
 * {@link TimestampQuery}).
 */
export type ComputePassDescriptor = GPUComputePassDescriptor;

// ─── Pipelines ───────────────────────────────────────────────────────────────

/**
 * WGSL sources for a render pipeline, in one of two mutually exclusive forms:
 * `shader`, a single source declaring both stages, or `vertex` with an optional
 * `fragment`. Omitting `fragment` from the split form builds a pipeline with no
 * fragment stage and no color targets — depth-only passes: shadow maps, depth
 * prepasses.
 */
export type RenderPipelineShaders =
  | { shader: string; vertex?: never; fragment?: never }
  | { shader?: never; vertex: string; fragment?: string };

interface RenderPipelineState
  extends
    GPUObjectDescriptorBase,
    Pick<GPUPrimitiveState, "topology" | "cullMode" | "frontFace">,
    Pick<
      GPUDepthStencilState,
      | "depthCompare"
      | "stencilFront"
      | "stencilBack"
      | "stencilReadMask"
      | "stencilWriteMask"
    > {
  /** Entry point. Defaults to the first `@vertex` function in the source. */
  vertexEntryPoint?: string;
  /** Entry point. Defaults to the first `@fragment` function in the source. */
  fragmentEntryPoint?: string;
  /**
   * Enables depth: defaults to false; when true, depthCompare defaults to
   * "less-equal" instead of "always".
   */
  depthWriteEnabled?: boolean;
  /** Blend state, per color target or one applied to all. */
  blend?: GPUBlendState | (GPUBlendState | undefined)[];
  /** Color write mask, per color target or one applied to all. */
  writeMask?: GPUColorWriteFlags | GPUColorWriteFlags[];
  /**
   * WGSL pipeline-overridable constant values, applied to both stages. Changing
   * them yields a distinct cached pipeline variant.
   */
  constants?: Record<string, number>;
  /** Constant depth bias added to fragments (eg. slope-scaled shadow bias). */
  depthBias?: number;
  /** Depth bias scaled by the fragment's depth slope. */
  depthBiasSlopeScale?: number;
  /** Maximum depth bias magnitude. */
  depthBiasClamp?: number;
  /**
   * Derive the sample coverage mask from the fragment's alpha at
   * `@location(0)`, instead of alpha being purely a colour channel.
   *
   * Antialiases cutouts — foliage, decals — which `discard` cannot: discard is
   * all-or-nothing per fragment, so a masked edge stays hard however many
   * samples the attachment has.
   *
   * Ignored when the pass is single-sampled, since there is no coverage mask to
   * derive. The fragment shader must not write `@builtin(sample_mask)`.
   */
  alphaToCoverage?: boolean;
}

/**
 * Lazy render pipeline description: a plain object, no GPU work at creation.
 * The GPURenderPipeline is created (and cached) at first submit, when vertex
 * layout and target formats are known.
 *
 * Shader sources are cached by object identity at first submit: don't swap
 * `shader`/`vertex`/`fragment` on an existing pipeline object — create a new
 * one instead. All other fields may change between submits.
 */
export type RenderPipeline = RenderPipelineState & RenderPipelineShaders;

export interface ComputePipeline
  extends
    GPUObjectDescriptorBase,
    Pick<GPUProgrammableStage, "entryPoint" | "constants"> {
  /**
   * WGSL compute shader source. Cached by object identity: don't swap the
   * source on an existing pipeline object.
   */
  compute: string;
}

// ─── Commands ────────────────────────────────────────────────────────────────

/** Values accepted in the uniforms dict, resolved by reflected binding kind. */
export type UniformValue =
  | PackableValue
  | Record<string, PackableValue>
  | GpuTexture
  | GPUTextureView
  | GPUSampler
  | GpuBuffer
  | HTMLVideoElement;

export type Uniforms = Record<string, UniformValue>;

export interface VertexAttribute {
  buffer: GpuBuffer;
  /** Byte offset into the buffer. Defaults to 0. */
  offset?: number;
  /**
   * Byte stride between elements, WebGPU's `GPUVertexBufferLayout.arrayStride`.
   * Defaults to the reflected WGSL type's tightly packed size.
   */
  arrayStride?: number;
  /** Defaults to "vertex". Set to "instance" for per-instance attributes. */
  stepMode?: GPUVertexStepMode;
  /**
   * Format override. Defaults to the format inferred from the shader's WGSL
   * type.
   */
  format?: GPUVertexFormat;
}

export type Attributes = Record<string, GpuBuffer | VertexAttribute>;

export interface RenderCommand extends GPUObjectDescriptorBase {
  pass?: RenderPassDescriptor;
  /** Omit for a pass-only command: an empty pass clears its attachments. */
  pipeline?: RenderPipeline;
  /** Vertex attributes keyed by the shader's input names. */
  attributes?: Attributes;
  indices?: GpuBuffer;
  uniforms?: Uniforms;
  /** Escape hatch: pre-built bind groups by group index, skipping reflection. */
  bindGroups?: Record<number, GPUBindGroup>;
  /**
   * VertexCount/indexCount. Inferred from indices or the first attribute when
   * omitted.
   */
  count?: number;
  instanceCount?: number;
  /** FirstVertex, or firstIndex when drawing with indices. */
  first?: number;
  baseVertex?: number;
  firstInstance?: number;
  /**
   * Draw parameters sourced from a GPU buffer —
   * drawIndirect/drawIndexedIndirect (replaces WebGL multi-draw).
   */
  indirect?: GpuBuffer | { buffer: GpuBuffer; offset?: number };
  /** SetViewport [x, y, width, height] — origin is top-left in WebGPU. */
  viewport?: [number, number, number, number];
  /** SetScissorRect [x, y, width, height]. */
  scissor?: [number, number, number, number];
  /** SetStencilReference: the value stencilFront/stencilBack compare against. */
  stencilReference?: GPUStencilValue;
}

export interface ComputeCommand extends GPUObjectDescriptorBase {
  /**
   * Pass options (label, timestampWrites). The pass itself is implied: one per
   * submit, shared by batch variants and pass scope submits.
   */
  pass?: ComputePassDescriptor;
  pipeline: ComputePipeline;
  uniforms?: Uniforms;
  bindGroups?: Record<number, GPUBindGroup>;
  /** DispatchWorkgroups counts, or a GPU buffer for dispatchWorkgroupsIndirect. */
  dispatch:
    | number
    | [number, number?, number?]
    | { buffer: GpuBuffer; offset?: number };
}

export interface BundleCommand extends GPUObjectDescriptorBase {
  pass?: RenderPassDescriptor;
  /** Replayed with executeBundles. */
  bundles: GPURenderBundle[];
}

export type Command = RenderCommand | ComputeCommand | BundleCommand;