import { createTimestampQuery } from "./query.js";
import { shared } from "./shared.js";
import type { TimestampQuery } from "./query.js";
import type { GpuContext } from "./types.js";

export type DebugLogLevel = "resources" | "commands";

export interface DebugOptions {
  /**
   * Measure GPU pass times and frame submit-to-completion latency, reported as
   * {@link DebugStats.timings}. Requires the "timestamp-query" feature on the
   * device — `optionalFeatures: ["timestamp-query"]` at createContext — and
   * stays off with a warning without it.
   */
  timing?: boolean;
  /**
   * "resources" logs GPU object creation, cache misses and allocator growth —
   * quiet in steady state, so per-frame messages indicate churn. "commands"
   * additionally logs every pass, draw and dispatch.
   */
  log?: DebugLogLevel | false;
  /**
   * Wrap frames and labeled commands in pushDebugGroup/popDebugGroup so GPU
   * profilers show the frame structure.
   */
  groups?: boolean;
}

const RESOURCE_COUNTS_KEYS = [
  "buffers",
  "textures",
  "querySets",
  "samplers",
  "renderBundles",
] as const;
const CACHE_COUNTS_KEYS = [
  "shaderReflections",
  "shaderModules",
  "renderPipelines",
  "computePipelines",
  "bindGroupLayouts",
  "pipelineLayouts",
  "bindGroups",
] as const;

/** Creation/disposal counts and estimated GPU memory for one resource type. */
export interface ResourceCounts {
  alive: number;
  created: number;
  disposed: number;
  /** Estimated GPU memory in bytes (texture sizes are format estimates). */
  bytes: number;
}
/**
 * Get-or-create cache hit/miss counts for one cache.
 *
 * An access is recorded by whichever layer served it: a draw reusing a pipeline
 * object's reflection or pipeline variant books the hit in the commands layer,
 * and never reaches the source- and descriptor-keyed caches under it — which
 * would otherwise only ever see the misses that got past the front line, and
 * report 0% on a scene where every draw is in fact cached.
 */
export interface CacheCounts {
  created: number;
  hits: number;
  /**
   * Wall time spent creating the missed entries, in milliseconds — the cost a
   * hit avoids. Only the miss path is measured, so a cache at steady state
   * leaves it untouched.
   */
  ms: number;
}
export type ResourceCountsKey = (typeof RESOURCE_COUNTS_KEYS)[number];
export type CacheCountsKey = (typeof CACHE_COUNTS_KEYS)[number];

/**
 * One live GPU resource. Frozen at creation and shared with every snapshot
 * rather than copied into each: WebGPU has no resize or reformat, so none of
 * these can go stale, and `bytes` is what untrackResource decrements the
 * counters by.
 */
export interface ResourceEntry {
  /** Same id the resource object carries, or an assigned one for raw handles. */
  readonly id: number;
  readonly label?: string;
  /** Estimated GPU memory in bytes; 0 for resources owning none. */
  readonly bytes: number;
}
export interface BufferEntry extends ResourceEntry {
  readonly usage: GPUBufferUsageFlags;
}
export interface TextureEntry extends ResourceEntry {
  readonly usage: GPUTextureUsageFlags;
  readonly format: GPUTextureFormat;
  readonly width: number;
  readonly height: number;
  readonly depthOrArrayLayers: number;
  readonly mipLevelCount: number;
  readonly sampleCount: number;
  readonly dimension: GPUTextureDimension;
}
export interface QuerySetEntry extends ResourceEntry {
  readonly type: GPUQueryType;
  readonly count: number;
}

/**
 * Every live resource by type, in creation order. Group it to break usage,
 * format or size down however the question demands:
 *
 * ```js
 * import { usageNames } from "pex-gpu/internals";
 *
 * Object.groupBy(resources.textures, (texture) => texture.format);
 * Object.groupBy(resources.buffers, (buffer) =>
 *   usageNames(buffer.usage, "buffer").join("|"),
 * );
 * ```
 */
export interface DebugResources {
  buffers: BufferEntry[];
  textures: TextureEntry[];
  querySets: QuerySetEntry[];
  samplers: ResourceEntry[];
  renderBundles: ResourceEntry[];
}

type UsageNames = readonly (readonly [number, string])[];

// Spec bit values rather than reads of GPUBufferUsage/GPUTextureUsage: this
// module is a leaf imported by every other one, and those globals are absent
// outside a WebGPU host (reflection/packing run in plain Node).
const BUFFER_USAGE_NAMES: UsageNames = [
  [0x00_01, "map-read"],
  [0x00_02, "map-write"],
  [0x00_04, "copy-src"],
  [0x00_08, "copy-dst"],
  [0x00_10, "index"],
  [0x00_20, "vertex"],
  [0x00_40, "uniform"],
  [0x00_80, "storage"],
  [0x01_00, "indirect"],
  [0x02_00, "query-resolve"],
];
const TEXTURE_USAGE_NAMES: UsageNames = [
  [0x01, "copy-src"],
  [0x02, "copy-dst"],
  [0x04, "texture-binding"],
  [0x08, "storage-binding"],
  [0x10, "render-attachment"],
];

/**
 * Split raw usage flags into WebGPU's flag names, in bit order — for grouping
 * or printing {@link DebugResources} entries.
 *
 * ```js
 * usageNames(positions.usage, "buffer").join("|"); // "copy-dst|vertex"
 * ```
 */
export const usageNames = (
  usage: number,
  type: "buffer" | "texture",
): string[] =>
  (type === "buffer" ? BUFFER_USAGE_NAMES : TEXTURE_USAGE_NAMES)
    .filter(([bit]) => (usage & bit) !== 0)
    .map(([, name]) => name);

const createResourceCounts = (): ResourceCounts => ({
  alive: 0,
  created: 0,
  disposed: 0,
  bytes: 0,
});
const createCacheCounts = (): CacheCounts => ({
  created: 0,
  hits: 0,
  ms: 0,
});

const createResourceCountsRecord = (): Record<
  ResourceCountsKey,
  ResourceCounts
> =>
  Object.fromEntries(
    RESOURCE_COUNTS_KEYS.map((key) => [key, createResourceCounts()]),
  ) as Record<ResourceCountsKey, ResourceCounts>;

const createCacheCountsRecord = (): Record<CacheCountsKey, CacheCounts> =>
  Object.fromEntries(
    CACHE_COUNTS_KEYS.map((key) => [key, createCacheCounts()]),
  ) as Record<CacheCountsKey, CacheCounts>;

/** Work submitted through declarative commands since the frame started. */
export interface FrameCounts extends Record<CacheCountsKey, CacheCounts> {
  renderPasses: number;
  computePasses: number;
  draws: number;
  dispatches: number;
  /** Bundles replayed with executeBundles. */
  bundles: number;
  /** Bytes packed into the uniform ring allocator. */
  uniformBytes: number;
}

/**
 * Resource/cache pressure counters since context creation — always on, never
 * reset.
 */
export interface LifetimeCounts
  extends
    Record<ResourceCountsKey, ResourceCounts>,
    Record<CacheCountsKey, CacheCounts> {}

/** The always-on counters, without the resource listing. */
export interface DebugCounters {
  lifetime: LifetimeCounts;
  frame: FrameCounts;
}

/**
 * Resource/rendering pressure for a context. `lifetime` is always on and starts
 * at context creation; `frame` resets at the start of each frame() tick;
 * `resources` lists what is alive right now; `timings` reports the GPU.
 */
export interface DebugStats extends DebugCounters {
  resources: DebugResources;
  /**
   * GPU timings for the most recently completed frame, or null while timing is
   * off or no readback has landed yet. A sibling of `frame` rather than a field
   * in it: those counters reset synchronously at beginFrame, and these describe
   * a frame two or three back — {@link DebugTimings.frameIndex} says which.
   */
  timings: DebugTimings | null;
}

type LiveResources = {
  [K in ResourceCountsKey]: Map<number, DebugResources[K][number]>;
};

const createLiveResources = (): LiveResources =>
  Object.fromEntries(
    RESOURCE_COUNTS_KEYS.map((key) => [key, new Map()]),
  ) as LiveResources;

interface DebugState {
  counters: DebugCounters;
  live: LiveResources;
  log: DebugLogLevel | false;
  groups: boolean;
  timing: boolean;
}

// Keyed by device rather than context so device-holding internals
// (PipelineCache, RingAllocator) can report without a context reference.
const states = shared("debugState", () => new WeakMap<GPUDevice, DebugState>());

const state = (device: GPUDevice): DebugState =>
  states.getOrInsertComputed(device, () => ({
    counters: {
      lifetime: {
        ...createResourceCountsRecord(),
        ...createCacheCountsRecord(),
      },
      frame: {
        renderPasses: 0,
        computePasses: 0,
        draws: 0,
        dispatches: 0,
        bundles: 0,
        uniformBytes: 0,
        ...createCacheCountsRecord(),
      },
    },
    live: createLiveResources(),
    log: false,
    groups: false,
    timing: false,
  }));

/**
 * Live mutable counters for a device — incremented by the instrumented
 * internals.
 */
export const debugCounters = (device: GPUDevice): DebugCounters =>
  state(device).counters;

/**
 * List a created resource and bump its type's lifetime counters. Takes
 * ownership of `entry` and freezes it. Entries hold plain data, never the GPU
 * handle: that would pin an undisposed resource that is otherwise unreachable,
 * turning a recoverable leak into a permanent one.
 */
export function trackResource<K extends ResourceCountsKey>(
  device: GPUDevice,
  type: K,
  entry: DebugResources[K][number],
): void {
  const { counters, live } = state(device);
  const counts = counters.lifetime[type];
  counts.created++;
  counts.alive++;
  counts.bytes += entry.bytes;
  Object.freeze(entry);
  live[type].set(entry.id, entry);
}

/** Drop a disposed resource. A second call for the same id is a no-op. */
export function untrackResource(
  device: GPUDevice,
  type: ResourceCountsKey,
  id: number,
): void {
  const { counters, live } = state(device);
  const entry = live[type].get(id);
  if (!entry) return;
  live[type].delete(id);
  const counts = counters.lifetime[type];
  counts.alive--;
  counts.disposed++;
  counts.bytes -= entry.bytes;
}

// The device is held weakly: a pending finalizer must not be what keeps a
// GPUDevice — and through it every context's state — alive.
const collected = new FinalizationRegistry<{
  device: WeakRef<GPUDevice>;
  type: ResourceCountsKey;
  id: number;
}>(({ device, type, id }) => {
  const value = device.deref();
  if (value) untrackResource(value, type, id);
});

/**
 * Like {@link trackResource}, for resources WebGPU reclaims with the JS object
 * instead of a destroy() — samplers and render bundles, which own no dispose()
 * to untrack them. Entries drop once `object` is collected.
 */
export function trackCollectable<K extends ResourceCountsKey>(
  device: GPUDevice,
  type: K,
  object: object,
  entry: DebugResources[K][number],
): void {
  trackResource(device, type, entry);
  collected.register(object, {
    device: new WeakRef(device),
    type,
    id: entry.id,
  });
}

/**
 * Record a get-or-create cache lookup against both the lifetime and per-frame
 * counters. `ms` is how long the miss took to produce its entry; null records a
 * hit, which by definition produced nothing.
 */
export function recordCacheAccess(
  device: GPUDevice,
  key: CacheCountsKey,
  ms: number | null,
): void {
  const counters = debugCounters(device);
  const lifetime = counters.lifetime[key];
  const frame = counters.frame[key];
  if (ms === null) {
    lifetime.hits++;
    frame.hits++;
  } else {
    lifetime.created++;
    frame.created++;
    lifetime.ms += ms;
    frame.ms += ms;
  }
}

const resetCacheCounts = (cache: CacheCounts): void => {
  cache.created = 0;
  cache.hits = 0;
  cache.ms = 0;
};

/** Reset the per-frame counters; called by frame() at the start of each tick. */
export function resetFrameCounters(device: GPUDevice): void {
  const frame = state(device).counters.frame;
  frame.renderPasses = 0;
  frame.computePasses = 0;
  frame.draws = 0;
  frame.dispatches = 0;
  frame.bundles = 0;
  frame.uniformBytes = 0;
  for (const key of CACHE_COUNTS_KEYS) resetCacheCounts(frame[key]);
}

/** Whether GPU debug groups should be emitted, per {@link debug}. */
export const debugGroupsEnabled = (device: GPUDevice): boolean =>
  states.get(device)?.groups ?? false;

/** Whether pass timing should be collected, per {@link debug}. */
export const timingEnabled = (device: GPUDevice): boolean =>
  states.get(device)?.timing ?? false;

/**
 * Log a debug message when the level is enabled. The message is built lazily so
 * disabled logging costs a WeakMap lookup.
 */
export function debugLog(
  device: GPUDevice,
  level: DebugLogLevel,
  message: () => string,
): void {
  const log = states.get(device)?.log;
  if (log === "commands" || log === level) {
    console.debug(`pex-gpu: ${message()}`);
  }
}

/**
 * Toggle debug logging and GPU debug groups. Counters ({@link debugStats}) are
 * always on — this only controls reporting.
 *
 * Timing is absent from the boolean shorthand on purpose: it needs a device
 * feature and changes what the GPU records, so it is only ever switched on by
 * asking for it.
 *
 * ```js
 * gpu.debug(ctx); // log: "resources", groups: true
 * gpu.debug(ctx, { log: "commands" }); // full submit stream, no groups
 * gpu.debug(ctx, { timing: true }); // pass timings, no logging
 * gpu.debug(ctx, false); // all reporting off
 * ```
 */
export function debug(
  ctx: GpuContext,
  options: boolean | DebugOptions = true,
): void {
  const value = state(ctx.device);
  if (typeof options === "boolean") {
    value.log = options && "resources";
    value.groups = options;
    value.timing = false;
  } else {
    value.log = options.log ?? false;
    value.groups = options.groups ?? false;
    const timing = options.timing ?? false;
    if (timing && !ctx.device.features.has("timestamp-query")) {
      console.warn(
        'pex-gpu: debug timing needs the "timestamp-query" feature — pass optionalFeatures: ["timestamp-query"] to createContext',
      );
    }
    value.timing = timing && ctx.device.features.has("timestamp-query");
  }
}

/**
 * Snapshot of the context's resource/rendering pressure counters and of every
 * live resource.
 *
 * ```js
 * const { lifetime, frame, resources } = gpu.debugStats(ctx);
 * console.log(`${lifetime.buffers.alive} buffers, ${frame.draws} draws`);
 * console.table(resources.textures);
 * ```
 */
export function debugStats(ctx: GpuContext): DebugStats {
  const { counters, live } = state(ctx.device);
  return {
    ...structuredClone(counters),
    resources: Object.fromEntries(
      RESOURCE_COUNTS_KEYS.map((key) => [key, [...live[key].values()]]),
    ) as unknown as DebugResources,
    // Rebuilt per readback, so the snapshot can share it like a resource entry.
    timings: timingStates.get(ctx)?.last ?? null,
  };
}

// ─── GPU timing ──────────────────────────────────────────────────────────────

/** One timed pass, in the order its encoder was opened. */
export interface PassTiming {
  index: number;
  label: string;
  kind: "render" | "compute";
  /**
   * GPU time in milliseconds, or null when the pass went untimed — it carried
   * caller-supplied `timestampWrites`, or the query set was full.
   */
  ms: number | null;
}

/**
 * Timings for one completed frame, reached through {@link DebugStats.timings}.
 * They lag the current frame by two or three, which is why they sit beside the
 * per-frame counters rather than in them.
 */
export interface DebugTimings {
  /** The frame these came from, counted from the first timed frame. */
  frameIndex: number;
  /**
   * Submit to queue completion, in milliseconds. Frame latency, not GPU busy
   * time: with frames in flight it includes waiting behind earlier ones.
   */
  latency: number | null;
  /**
   * Sum of the pass times. An attribution rather than an exclusive cost —
   * passes can overlap, and a pass's span includes stalls waiting on its own
   * dependencies, so this can exceed {@link DebugTimings.latency}.
   */
  gpuTime: number | null;
  passes: PassTiming[];
  /**
   * Every timing is a multiple of 100µs, so the browser is bucketing them.
   * Chrome does by default: run with
   * chrome://flags/#enable-webgpu-developer-features (or
   * --disable-dawn-features=timestamp_quantization) for real numbers.
   */
  quantized: boolean;
  /** Frames skipped because every readback buffer was still in flight. */
  dropped: number;
}

interface PendingPass {
  label: string;
  kind: "render" | "compute";
  /** Index of the begin query, or null when the pass went untimed. */
  query: number | null;
}

/** A resolved frame, waiting for its timestamps and its latency to land. */
interface ArmedFrame {
  frameIndex: number;
  passes: PendingPass[];
  latency: number | null;
}

interface TimingState {
  query: TimestampQuery | null;
  /** Query slots the next query set is created with. */
  capacity: number;
  frameIndex: number;
  /** Queries claimed this frame. */
  cursor: number;
  /** Queries this frame asked for, overflow included — drives growth. */
  needed: number;
  passes: PendingPass[];
  armed: ArmedFrame | null;
  dropped: number;
  last: DebugTimings | null;
}

/** Two queries per pass; the spec caps a query set at 4096. */
const INITIAL_QUERIES = 256;
const MAX_QUERIES = 4096;
/**
 * Readback buffers. Two frames may be in flight
 * (`FrameOptions.maxFramesInFlight`) and a third is mapping, so three is what
 * it takes to time every frame without ever waiting on one.
 */
const BUFFERING = 3;
/** Chrome's default timestamp bucket, in nanoseconds. */
const QUANTUM = 100_000n;

const timingStates = shared(
  "timingState",
  () => new WeakMap<GpuContext, TimingState>(),
);

const timingState = (ctx: GpuContext): TimingState =>
  timingStates.getOrInsertComputed(ctx, () => ({
    query: null,
    capacity: INITIAL_QUERIES,
    frameIndex: -1,
    cursor: 0,
    needed: 0,
    passes: [],
    armed: null,
    dropped: 0,
    last: null,
  }));

/**
 * Start a frame's timings: called by beginFrame. Grows the query set when the
 * previous frame ran out of it — a set can't be resized, so the passes that
 * overflowed went untimed and the next frame gets the bigger one.
 */
export function resetFrameTiming(ctx: GpuContext): void {
  if (!timingEnabled(ctx.device)) return;
  const timing = timingState(ctx);

  if (timing.needed > timing.capacity) {
    const capacity = Math.min(
      MAX_QUERIES,
      2 ** Math.ceil(Math.log2(timing.needed)),
    );
    if (capacity > timing.capacity) {
      timing.capacity = capacity;
      // The replaced query self-destructs once its last readback lands.
      timing.query?.dispose();
      timing.query = null;
      debugLog(
        ctx.device,
        "resources",
        () => `pass timing query set grown to ${capacity} queries`,
      );
    }
  }

  timing.frameIndex++;
  timing.cursor = 0;
  timing.needed = 0;
  timing.passes = [];
}

/**
 * Claim a begin/end query pair for a pass being opened, and return the writes
 * to put in its descriptor. Undefined when timing is off, when `userWrites`
 * already claims the descriptor's single `timestampWrites` slot, or when the
 * query set is full — the pass is still listed, with a null time.
 *
 * Render and compute timestamp writes are the same shape, so one return type
 * serves both descriptors.
 */
export function trackPassTiming(
  ctx: GpuContext,
  label: string,
  kind: "render" | "compute",
  userWrites?: GPURenderPassTimestampWrites | GPUComputePassTimestampWrites,
): GPURenderPassTimestampWrites | undefined {
  if (!timingEnabled(ctx.device)) return;
  const timing = timingState(ctx);
  if (userWrites) {
    timing.passes.push({ label, kind, query: null });
    return;
  }
  // Counted before the capacity check, and only for passes that would use the
  // set, so overflow is what grows it and caller-timed passes never do.
  timing.needed += 2;
  const query = (timing.query ??= createTimestampQuery(ctx, {
    label: "pex-gpu pass timings",
    count: timing.capacity,
    buffering: BUFFERING,
  }));
  if (timing.cursor + 2 > query.count) {
    timing.passes.push({ label, kind, query: null });
    return;
  }
  const begin = timing.cursor;
  timing.cursor += 2;
  timing.passes.push({ label, kind, query: begin });
  return {
    querySet: query.querySet,
    beginningOfPassWriteIndex: begin,
    endOfPassWriteIndex: begin + 1,
  };
}

/**
 * Encode this frame's query resolution: called by endFrame, after every pass
 * has ended and before the encoder is finished.
 */
export function resolveFrameTiming(
  ctx: GpuContext,
  encoder: GPUCommandEncoder,
): void {
  if (!timingEnabled(ctx.device)) return;
  const timing = timingState(ctx);
  if (!timing.query || timing.cursor === 0) return;

  if (!timing.query.resolve(encoder, timing.cursor)) {
    // Every readback buffer is still mapping: time the next frame instead.
    timing.dropped++;
    return;
  }
  timing.armed = {
    frameIndex: timing.frameIndex,
    passes: timing.passes,
    latency: null,
  };
}

function buildTimings(
  timing: TimingState,
  armed: ArmedFrame,
  timestamps: BigInt64Array,
): DebugTimings {
  let gpuTime = 0;
  let timed = 0;
  let quantized = true;
  const passes = armed.passes.map((pass, index): PassTiming => {
    let ms: number | null = null;
    if (pass.query !== null) {
      const begin = timestamps[pass.query]!;
      const end = timestamps[pass.query + 1]!;
      // A zero timestamp is one the GPU never wrote.
      if (begin > 0n && end > begin) {
        const ns = end - begin;
        quantized &&= ns % QUANTUM === 0n;
        ms = Number(ns) / 1e6;
        gpuTime += ms;
        timed++;
      }
    }
    return { index, label: pass.label, kind: pass.kind, ms };
  });

  return {
    frameIndex: armed.frameIndex,
    latency: armed.latency,
    gpuTime: timed ? gpuTime : null,
    passes,
    quantized: timed > 0 && quantized,
    dropped: timing.dropped,
  };
}

/**
 * Close a frame's timings: called by endFrame right after the queue submit,
 * which is both when the resolved buffer may be mapped and when frame latency
 * starts counting.
 */
export function completeFrameTiming(ctx: GpuContext): void {
  if (!timingEnabled(ctx.device)) return;
  const timing = timingState(ctx);
  const armed = timing.armed;
  timing.armed = null;
  if (!armed || !timing.query) return;

  const submittedAt = performance.now();
  void ctx.device.queue.onSubmittedWorkDone().then(
    () => {
      armed.latency = performance.now() - submittedAt;
    },
    () => {},
  );
  void timing.query.read().then((timestamps) => {
    if (timestamps) timing.last = buildTimings(timing, armed, timestamps);
  });
}

/** Release the query set and readback buffers. Called by context dispose. */
export function disposeTiming(ctx: GpuContext): void {
  const timing = timingStates.get(ctx);
  if (!timing) return;
  timing.query?.dispose();
  timingStates.delete(ctx);
}