import { PipelineCache } from "./pipeline-cache.js";
import {
  completeFrameTiming,
  debugGroupsEnabled,
  disposeTiming,
  resetFrameCounters,
  resetFrameTiming,
  resolveFrameTiming,
} from "./debug.js";
import { commandsState, peekCommandsState } from "./state.js";
import type { CommandsState } from "./state.js";
import { assertSchema, brand, shared } from "./shared.js";
import { createTexture } from "./texture.js";
import type {
  FrameInfo,
  FrameSegment,
  GpuContext,
  GpuTexture,
} from "./types.js";

export interface CreateContextOptions
  extends
    GPURequestAdapterOptions,
    Partial<Omit<GPUCanvasConfiguration, "device">>,
    Pick<GPUDeviceDescriptor, "requiredLimits" | "label" | "defaultQueue"> {
  /**
   * Existing canvas to render into. One is created and appended to body
   * otherwise.
   */
  canvas?: HTMLCanvasElement;
  /** CSS width. Defaults to the window size when creating a canvas. */
  width?: number;
  /** CSS height. */
  height?: number;
  /**
   * Physical pixels per CSS pixel. Defaults to 1. Bounded by the device's
   * maxTextureDimension2D; cost grows with its square.
   */
  pixelRatio?: number;
  /**
   * Features to require on the device, eg. "timestamp-query". Unlike
   * requestDevice (which rejects), missing ones are skipped with a warning.
   */
  requiredFeatures?: Iterable<GPUFeatureName>;
  /**
   * Features to take when the adapter offers them, skipped silently otherwise.
   *
   * For capabilities an app adapts to rather than depends on — notably the
   * compressed-texture families, where no single device has all three, so
   * requesting them as `requiredFeatures` would warn on every run. WebGPU only
   * reports a feature on the device if it was asked for at creation, so there
   * is no way to opt in later; read `device.features` afterwards to see what
   * was granted.
   */
  optionalFeatures?: Iterable<GPUFeatureName>;
  /**
   * Limits to take up to what the adapter supports, clamped instead of
   * rejected.
   *
   * `requiredLimits` reaches requestDevice untouched, so an unsatisfiable value
   * rejects — what a hard dependency wants. These are best effort instead:
   * unknown keys are dropped, `max*` values clamp down to the adapter's and
   * `min*` alignments clamp up. `requiredLimits` wins on conflict; read
   * `device.limits` afterwards for what was granted.
   */
  optionalLimits?: NonNullable<GPUDeviceDescriptor["requiredLimits"]>;
  /**
   * Attach a context-managed depth buffer to canvas-targeting passes. Defaults
   * to true.
   */
  depth?: boolean;
  /**
   * Include a stencil aspect in the canvas depth/stencil buffer. Defaults to
   * false.
   */
  stencil?: boolean;
  /**
   * Explicit format for the canvas depth/stencil buffer, eg. "depth32float".
   * Overrides `depth`/`stencil`.
   */
  depthStencilFormat?: GPUTextureFormat;
  /**
   * Keep the canvas's rendered content across frames instead of each frame
   * starting from undefined content (WebGPU gives canvas textures no such
   * guarantee — this restores the WebGL-era `preserveDrawingBuffer`
   * ergonomics). Implemented with a persistent backing texture copied into the
   * canvas after every frame, so it costs one extra texture-sized copy per
   * frame — leave it off unless a pass relies on `loadOp: "load"` carrying over
   * prior frames (eg. trail/feedback effects) or on reading back the canvas
   * without re-rendering it first. Defaults to false. Requires the canvas usage
   * to include `COPY_DST`, added automatically unless `usage` is set
   * explicitly.
   */
  preserveDrawingBuffer?: boolean;
}

function resolveDepthStencilFormat(
  options: CreateContextOptions,
): GPUTextureFormat | undefined {
  if (options.depthStencilFormat) return options.depthStencilFormat;
  const depth = options.depth ?? true;
  if (depth) return options.stencil ? "depth24plus-stencil8" : "depth24plus";
  return options.stencil ? "stencil8" : undefined;
}

interface ContextState {
  resized: boolean;
  running: boolean;
  preserveDrawingBuffer: boolean;
  /** Last ratio asked for, so a clamped resize doesn't degrade later ones. */
  requestedPixelRatio: number;
  /**
   * Persistent canvas-sized color texture backing `preserveDrawingBuffer`,
   * recreated on resize.
   */
  backbuffer: GpuTexture | null;
}

const contextState = shared(
  "contextState",
  () => new WeakMap<GpuContext, ContextState>(),
);

const getState = (ctx: GpuContext): ContextState => {
  const state = contextState.get(ctx);
  if (!state) {
    assertSchema(ctx);
    throw new Error("pex-gpu: unknown or disposed context");
  }
  return state;
};

/**
 * Create a WebGPU context: adapter, device, and a configured canvas. `label`
 * names the device in browser error messages and tooling; `format` defaults to
 * getPreferredCanvasFormat(), eg. "rgba16float" for HDR.
 *
 * ```js
 * const ctx = await gpu.createContext({ width: 640, height: 480 });
 * ```
 */
export async function createContext(
  options: CreateContextOptions = {},
): Promise<GpuContext> {
  if (!navigator.gpu) {
    throw new Error("pex-gpu: WebGPU is not available in this browser");
  }

  const adapter = await navigator.gpu.requestAdapter({
    ...(options.powerPreference && {
      powerPreference: options.powerPreference,
    }),
    ...(options.featureLevel && { featureLevel: options.featureLevel }),
    ...(options.forceFallbackAdapter && { forceFallbackAdapter: true }),
    ...(options.xrCompatible && { xrCompatible: true }),
  });
  if (!adapter) throw new Error("pex-gpu: no WebGPU adapter available");

  const requestedFeatures = [...(options.requiredFeatures ?? [])];
  const missingFeatures = requestedFeatures.filter(
    (feature) => !adapter.features.has(feature),
  );
  if (missingFeatures.length) {
    console.warn(
      `pex-gpu: unsupported features skipped: ${missingFeatures.join(", ")}`,
    );
  }

  // Deduplicated: a feature listed as both required and optional is one entry.
  const grantedFeatures = new Set(
    [...requestedFeatures, ...(options.optionalFeatures ?? [])].filter(
      (feature) => adapter.features.has(feature),
    ),
  );

  const supportedLimits = adapter.limits as unknown as Record<
    string,
    number | undefined
  >;
  // "min" limits are alignments — lower is better, and the value must stay a
  // power of two — everything else is a maximum. Unknown or non-finite entries
  // are dropped rather than passed on, since requestDevice rejects on those.
  const clampedLimits: Record<string, number> = {};
  for (const [name, value] of Object.entries(options.optionalLimits ?? {})) {
    const supported = supportedLimits[name];
    if (supported === undefined || value === undefined) continue;
    if (!Number.isFinite(value)) continue;
    clampedLimits[name] = name.startsWith("min")
      ? Math.max(2 ** Math.floor(Math.log2(value)), supported)
      : Math.min(value, supported);
  }
  const requiredLimits = { ...clampedLimits, ...options.requiredLimits };

  const device = await adapter.requestDevice({
    ...(options.label && { label: options.label }),
    requiredFeatures: [...grantedFeatures],
    ...(Object.keys(requiredLimits).length > 0 && { requiredLimits }),
    ...(options.defaultQueue && { defaultQueue: options.defaultQueue }),
  });

  void (async () => {
    const info = await device.lost;
    if (info.reason !== "destroyed") {
      console.error(`pex-gpu: device lost: ${info.message}`);
    }
  })();

  // Chrome logs uncaptured errors to the console by default, Safari does not.
  // Handling the event (+ preventDefault) makes reporting identical everywhere.
  device.addEventListener("uncapturederror", (event) => {
    event.preventDefault();
    console.error(
      `pex-gpu: uncaptured ${event.error.constructor.name}: ${event.error.message}`,
    );
  });

  const canvas = options.canvas ?? document.createElement("canvas");
  if (!options.canvas) document.body.append(canvas);

  const canvasContext = canvas.getContext("webgpu");
  if (!canvasContext)
    throw new Error("pex-gpu: could not get a webgpu canvas context");

  const format = options.format ?? navigator.gpu.getPreferredCanvasFormat();
  const alphaMode = options.alphaMode ?? "opaque";
  canvasContext.configure({
    device,
    format,
    alphaMode,
    ...(options.usage
      ? { usage: options.usage }
      : options.preserveDrawingBuffer && {
          usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_DST,
        }),
    ...(options.colorSpace && { colorSpace: options.colorSpace }),
    ...(options.toneMapping && { toneMapping: options.toneMapping }),
    ...(options.viewFormats && { viewFormats: options.viewFormats }),
  });

  const dispose = (): void => {
    const state = contextState.get(ctx);
    if (state) {
      state.running = false;
      state.backbuffer?.dispose();
    }
    contextState.delete(ctx);
    // Only tears down declarative-commands state if it was ever touched —
    // code that never submits a command shouldn't pay to clean up an
    // allocation it never created.
    const commands = peekCommandsState(ctx);
    if (commands) {
      commands.allocator?.dispose();
      commands.storageAllocator?.dispose();
      commands.canvasDepthStencil?.dispose();
      commands.bindGroupLayouts.clear();
      commands.pipelineLayouts.clear();
      commands.bindGroups.clear();
      commands.bindGroupsByResource.clear();
    }
    ctx.pipelineCache.clear();
    disposeTiming(ctx);
    canvasContext.unconfigure();
    device.destroy();
  };

  const ctx: GpuContext = {
    adapter,
    device,
    canvas,
    canvasContext,
    format,
    alphaMode,
    depthStencilFormat: resolveDepthStencilFormat(options),
    pipelineCache: new PipelineCache(device),
    pixelRatio: 1,
    get width() {
      return canvas.width;
    },
    get height() {
      return canvas.height;
    },
    dispose,
    [Symbol.dispose]: dispose,
  };

  brand(ctx);
  contextState.set(ctx, {
    resized: false,
    running: false,
    preserveDrawingBuffer: options.preserveDrawingBuffer ?? false,
    requestedPixelRatio: 1,
    backbuffer: null,
  });
  resize(
    ctx,
    options.width ?? (options.canvas ? canvas.clientWidth : window.innerWidth),
    options.height ??
      (options.canvas ? canvas.clientHeight : window.innerHeight),
    options.pixelRatio ?? 1,
  );

  console.info(`pex-gpu ✔`);

  return ctx;
}

/**
 * Set the canvas CSS size and physical drawing buffer size. Size-dependent
 * resources (depth textures, render targets) are the caller's to recreate — the
 * next frame reports `resized: true`.
 *
 * `ctx.pixelRatio` reports the ratio applied: the requested one, unless the
 * drawing buffer would exceed maxTextureDimension2D.
 */
export function resize(
  ctx: GpuContext,
  width: number,
  height: number,
  pixelRatio?: number,
): void {
  const state = getState(ctx);
  if (pixelRatio) state.requestedPixelRatio = pixelRatio;

  const maxDimension = ctx.device.limits.maxTextureDimension2D;
  const ratio = Math.min(
    state.requestedPixelRatio,
    maxDimension / Math.max(width, height),
  );
  if (ratio < state.requestedPixelRatio) {
    console.warn(
      `pex-gpu: ${width}x${height} at pixelRatio ${state.requestedPixelRatio} exceeds maxTextureDimension2D (${maxDimension}), using pixelRatio ${ratio}`,
    );
  }

  ctx.pixelRatio = ratio;
  ctx.canvas.style.width = `${width}px`;
  ctx.canvas.style.height = `${height}px`;
  ctx.canvas.width = Math.max(1, Math.floor(width * ratio));
  ctx.canvas.height = Math.max(1, Math.floor(height * ratio));

  state.resized = true;
  if (!state.preserveDrawingBuffer) return;

  state.backbuffer?.dispose();
  state.backbuffer = createTexture(ctx, {
    width: ctx.width,
    height: ctx.height,
    format: ctx.format,
    usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC,
    label: "pex-gpu canvas backbuffer",
  });
}

/**
 * Open a recording segment: creates a command encoder, acquires the canvas
 * texture and resets the per-frame uniform ring allocator. submit() and other
 * declarative commands are only valid between a beginFrame()/endFrame() pair.
 *
 * Unlike {@link frame}, this doesn't touch requestAnimationFrame — call it from
 * any driving loop (a manual rAF loop, a fixed-timestep export loop, an async
 * pump awaiting backpressure from something else). `frame()` is a thin
 * convenience wrapper implemented on top of this pair for the common
 * display-synced case; reach for beginFrame()/endFrame() directly when a
 * frame's timing isn't display-refresh-driven — eg. rendering for precise video
 * export, where every rendered frame must be captured regardless of vsync.
 *
 * Segments don't nest: call endFrame() before opening another one.
 *
 * ```js
 * const segment = gpu.beginFrame(ctx);
 * gpu.submit(ctx, drawCmd);
 * gpu.endFrame(ctx);
 * ```
 */
export function beginFrame(
  ctx: GpuContext,
  label = "pex-gpu frame",
): FrameSegment {
  const state = getState(ctx);
  const commands = commandsState(ctx);
  if (commands.frame) {
    throw new Error(
      "pex-gpu: a frame segment is already open — call endFrame() before beginFrame() again",
    );
  }

  resetFrameCounters(ctx.device);
  resetFrameTiming(ctx);
  const encoder = ctx.device.createCommandEncoder({ label });
  // The swapchain texture only ever receives the final blit in endFrame()
  // when preserveDrawingBuffer is on — the segment itself renders into the
  // persistent backbuffer so content survives across segments.
  const swapchainTexture = ctx.canvasContext.getCurrentTexture();
  const canvasTexture = state.backbuffer?.texture ?? swapchainTexture;
  const canvasView = state.backbuffer?.view ?? swapchainTexture.createView();
  if (debugGroupsEnabled(ctx.device)) encoder.pushDebugGroup(label);

  commands.allocator?.reset();
  commands.storageAllocator?.reset();
  commands.frame = {
    encoder,
    swapchainTexture,
    canvasView,
    canvasDepthStencilCleared: false,
    activePass: null,
    activeComputePass: null,
  };

  return {
    encoder,
    canvasTexture,
    canvasView,
    width: ctx.width,
    height: ctx.height,
  };
}

// Uniforms staged by submit() must reach the GPU buffer before the segment
// encoder is submitted.
const uploadUniforms = (commands: CommandsState) => {
  commands.allocator?.flush();
  commands.storageAllocator?.flush();
};

/**
 * Submit the open segment's commands so far and continue it in a new command
 * buffer. A validation error fails only the command buffer it is in, so work
 * flushed this way survives anything invalid recorded after it.
 *
 * Not valid inside a pass scope. Returns the new encoder: one obtained before
 * the flush, eg. `FrameSegment.encoder`, is finished.
 */
export function flush(ctx: GpuContext): GPUCommandEncoder {
  const commands = commandsState(ctx);
  const frame = commands.frame;
  if (!frame) {
    throw new Error("pex-gpu: flush() called without an open frame segment");
  }
  if (frame.activePass || frame.activeComputePass) {
    throw new Error("pex-gpu: flush() cannot be called inside a pass scope");
  }

  uploadUniforms(commands);
  const { label } = frame.encoder;
  const debugGroups = debugGroupsEnabled(ctx.device);
  if (debugGroups) frame.encoder.popDebugGroup();
  ctx.device.queue.submit([frame.encoder.finish()]);
  frame.encoder = ctx.device.createCommandEncoder({ label });
  if (debugGroups) frame.encoder.pushDebugGroup(label);
  return frame.encoder;
}

/**
 * Close the segment opened by {@link beginFrame}: ends any pass left open,
 * flushes the uniform ring allocator, blits the preserveDrawingBuffer
 * backbuffer if enabled, and submits the encoder. Always submits — there's no
 * discard path, so every opened segment's work reaches the GPU.
 */
export function endFrame(ctx: GpuContext): void {
  const state = getState(ctx);
  const commands = commandsState(ctx);
  const frame = commands.frame;
  if (!frame) {
    throw new Error(
      "pex-gpu: endFrame() called without a matching beginFrame()",
    );
  }
  commands.frame = null;

  if (frame.activePass) frame.activePass.encoder.end();
  if (frame.activeComputePass) frame.activeComputePass.encoder.end();
  uploadUniforms(commands);
  // Blit the persistent backbuffer into this segment's actual swapchain
  // texture so it's what gets presented, then the backbuffer itself is left
  // untouched for the next segment to build on.
  if (state.backbuffer) {
    frame.encoder.copyTextureToTexture(
      { texture: state.backbuffer.texture },
      { texture: frame.swapchainTexture },
      { width: ctx.width, height: ctx.height },
    );
  }

  resolveFrameTiming(ctx, frame.encoder);
  if (debugGroupsEnabled(ctx.device)) frame.encoder.popDebugGroup();
  ctx.device.queue.submit([frame.encoder.finish()]);
  // After the submit: the resolved buffer can only be mapped once the work
  // using it is queued, and frame latency counts from the same moment.
  completeFrameTiming(ctx);
}

export interface FrameOptions {
  /**
   * Frames allowed to be submitted but not yet finished before the loop waits.
   * Default 2.
   *
   * RequestAnimationFrame fires at display rate whether or not the GPU is
   * keeping up, so a frame costing more than the display interval would
   * otherwise queue work faster than it drains, without bound. That does not
   * degrade to a low frame rate — it degrades to the browser killing its GPU
   * process for not responding. Waiting on an older frame turns it back into
   * honest slowness.
   *
   * `Infinity` restores the unbounded behaviour, for a loop that wants to
   * measure or manage latency itself.
   */
  maxFramesInFlight?: number;
}

/**
 * Render loop. Each frame is a {@link beginFrame}/{@link endFrame} segment synced
 * to requestAnimationFrame. Return `false` from the callback to stop the loop —
 * the stopping frame's work is still submitted.
 *
 * Paced: at most {@link FrameOptions.maxFramesInFlight} frames may be in flight,
 * so a scene the GPU cannot keep up with drops frames instead of accumulating
 * submissions.
 *
 * The callback can drive the encoder directly, submit() declarative commands,
 * or mix both.
 *
 * The callback may be async: the segment stays open until it settles and the
 * next frame is only requested then, so two ticks never share an encoder and
 * everything drawn still lands in one command buffer. Await CPU work only — the
 * swapchain texture is acquired when the segment opens, so yielding long enough
 * for the browser to present would invalidate it.
 *
 * ```js
 * gpu.frame(ctx, ({ time }) => {
 *   gpu.submit(ctx, drawCmd);
 * });
 *
 * gpu.frame(ctx, ({ encoder, canvasView }) => {
 *   const pass = encoder.beginRenderPass({ colorAttachments: [{ view: canvasView, ... }] });
 *   // ...
 *   pass.end();
 * });
 *
 * gpu.frame(ctx, async () => {
 *   await renderer.render(); // one segment, one submit
 *   gui.draw();
 * });
 * ```
 */
export function frame(
  ctx: GpuContext,
  callback: (info: FrameInfo) => void | boolean | Promise<void | boolean>,
  options: FrameOptions = {},
): void {
  const state = getState(ctx);
  state.running = true;
  state.resized = false; // setup-time resize is not a runtime resize
  let startTime = -1;
  let previousTimestamp = -1;
  let frameIndex = 0;

  // At least one, or the wait below would have nothing to wait on and spin.
  const maxFramesInFlight = Math.max(1, options.maxFramesInFlight ?? 2);
  const paced = Number.isFinite(maxFramesInFlight);
  // One entry per frame submitted and not yet known to have finished.
  // onSubmittedWorkDone() covers everything submitted before it, so the entry
  // recorded `maxFramesInFlight` frames ago is exactly the one to wait on.
  const inFlight: Promise<void>[] = [];

  const tick = async (timestamp: number) => {
    if (!state.running || !contextState.has(ctx)) return;

    // Before the segment opens, never inside it: the swapchain texture is
    // acquired by beginFrame, and waiting here is long enough for the browser
    // to present, which would invalidate it.
    while (inFlight.length >= maxFramesInFlight) {
      await inFlight.shift();
      if (!state.running || !contextState.has(ctx)) return;
    }

    if (startTime === -1) startTime = previousTimestamp = timestamp;

    const time = (timestamp - startTime) / 1000;
    const deltaTime = (timestamp - previousTimestamp) / 1000;
    previousTimestamp = timestamp;

    const segment = beginFrame(ctx, `pex-gpu frame ${frameIndex}`);
    const info: FrameInfo = {
      ...segment,
      time,
      deltaTime,
      frame: frameIndex++,
      timestamp,
      resized: state.resized,
    };
    state.resized = false;

    let result: void | boolean;
    try {
      result = await callback(info);
    } finally {
      endFrame(ctx);
    }

    // Rejections are swallowed deliberately: this promise is a pacing signal,
    // not a result. A lost device rejects it, and a loop that stops pacing is a
    // better failure than one that never runs again.
    // Skipped when unpaced: nothing would ever drain the list.
    if (paced) {
      inFlight.push(ctx.device.queue.onSubmittedWorkDone().catch(() => {}));
    }

    if (result === false) {
      state.running = false;
      return;
    }

    requestAnimationFrame(tick);
  };

  requestAnimationFrame(tick);
}