import { debugCounters, debugLog, recordCacheAccess } from "./debug.js";
import { parseWGSL } from "./reflection.js";
import type { RenderPipelineShaders, ShaderReflection } from "./types.js";

/**
 * Unfold either shader form into per-stage sources: a combined `shader` becomes
 * the same string on both stages, so the two spellings share cache entries. An
 * undefined fragment means a pipeline with no fragment stage.
 */
export const resolveShaders = (
  shaders: RenderPipelineShaders,
): { vertex: string; fragment: string | undefined } => ({
  vertex: shaders.shader ?? shaders.vertex,
  fragment: shaders.shader ?? shaders.fragment,
});

interface RenderPipelineSourceState
  extends
    GPUObjectDescriptorBase,
    Partial<
      Pick<
        GPURenderPipelineDescriptor,
        "layout" | "primitive" | "depthStencil" | "multisample"
      >
    > {
  /**
   * Entry points default to the first `@vertex`/`@fragment` function in the
   * source.
   */
  vertexEntryPoint?: string;
  fragmentEntryPoint?: string;
  buffers?: GPUVertexBufferLayout[];
  targets?: GPUColorTargetState[];
  /** WGSL pipeline-overridable constants, applied to both stages. */
  constants?: GPUProgrammableStage["constants"];
  /** A stable cache key, required alongside a non-"auto" layout. */
  layoutKey?: string;
}

/**
 * Render pipeline description with WGSL sources instead of shader modules, so
 * descriptors stay serializable and cacheable. `layout` defaults to "auto".
 */
export type RenderPipelineSource = RenderPipelineSourceState &
  RenderPipelineShaders;

/** Compute counterpart of {@link RenderPipelineSource}. */
export interface ComputePipelineSource
  extends
    GPUObjectDescriptorBase,
    Pick<GPUProgrammableStage, "entryPoint" | "constants">,
    Partial<Pick<GPUComputePipelineDescriptor, "layout">> {
  compute: string;
  /** A stable cache key, required alongside a non-"auto" layout. */
  layoutKey?: string;
}

/**
 * Narrow a stage's constants to `referenced` — the overrides its entry point
 * actually uses (from reflection) — so a stage never receives a value for an
 * override it doesn't use: fatal on WebKit (a combined vertex+fragment source
 * declares both stages' overrides), ignored by Chrome. `referenced` is
 * undefined only when the entry point couldn't be resolved, in which case
 * constants pass through unscoped. Undefined result when nothing remains, so
 * `constants` is omitted entirely.
 */
function scopeConstants(
  constants: GPUProgrammableStage["constants"],
  referenced: ReadonlySet<string> | undefined,
): GPUProgrammableStage["constants"] | undefined {
  if (!constants) return undefined;
  if (!referenced) return constants;
  const scoped: Record<string, GPUPipelineConstantValue> = {};
  for (const [key, value] of Object.entries(constants)) {
    if (referenced.has(key)) scoped[key] = value;
  }
  return Object.keys(scoped).length ? scoped : undefined;
}

const formatCompilationMessage = (
  code: string,
  message: GPUCompilationMessage,
): string => {
  // lineNum/linePos are 1-based; 0 means the message has no source location
  const location = message.lineNum
    ? ` at ${message.lineNum}:${message.linePos}`
    : "";
  let out = `${message.type}${location}: ${message.message}`;
  const line = message.lineNum && code.split("\n")[message.lineNum - 1];
  if (line) {
    out += `\n  ${line}\n  ${" ".repeat(Math.max(0, message.linePos - 1))}^`;
  }
  return out;
};

async function reportShaderDiagnostics(
  code: string,
  module: GPUShaderModule,
  scopeError: Promise<GPUError | null>,
  label?: string,
): Promise<void> {
  // Each result is collected independently: a broken/missing
  // getCompilationInfo (Safari) must never mask the captured scope error,
  // and vice versa.
  let error: GPUError | null = null;
  try {
    error = await scopeError;
  } catch {
    // device loss: nothing left to report
  }
  let messages: readonly GPUCompilationMessage[] = [];
  try {
    ({ messages } = await module.getCompilationInfo());
  } catch {
    // getCompilationInfo unsupported or broken: the scope error still reports
  }
  if (!messages.length && !error) return;

  const name = label ? ` in "${label}"` : "";
  const warnings = messages.filter((message) => message.type !== "error");
  if (warnings.length) {
    console.warn(
      `pex-gpu: WGSL warnings${name}\n${warnings
        .map((message) => formatCompilationMessage(code, message))
        .join("\n")}`,
    );
  }

  const errors = messages.filter((message) => message.type === "error");
  if (!(errors.length || error)) return;

  const details = errors.length
    ? errors
        .map((message) => formatCompilationMessage(code, message))
        .join("\n")
    : error!.message;
  console.error(`pex-gpu: WGSL compilation failed${name}\n${details}`);
}

/** A pipeline stage's source + module, for post-failure diagnostics. */
interface StageModule {
  stage: string;
  code: string;
  module: GPUShaderModule;
}

/** Resolves to whether the pipeline failed validation. */
async function reportPipelineError(
  scopeError: Promise<GPUError | null>,
  label: string,
  stages: StageModule[],
): Promise<boolean> {
  let error: GPUError | null = null;
  try {
    error = await scopeError;
  } catch {
    return false; // device loss: nothing left to report
  }
  if (!error) return false;

  const details: string[] = [];
  for (const { stage, code, module } of stages) {
    let messages: readonly GPUCompilationMessage[] = [];
    try {
      ({ messages } = await module.getCompilationInfo());
    } catch {
      continue; // getCompilationInfo unsupported or broken on this browser
    }
    for (const message of messages) {
      details.push(`[${stage}] ${formatCompilationMessage(code, message)}`);
    }
  }

  console.error(
    `pex-gpu: pipeline creation failed in "${label}"\n${error.message}${
      details.length ? `\n\n${details.join("\n")}` : ""
    }`,
  );
  return true;
}

/**
 * Caches shader modules by source and pipelines by their full descriptor, so
 * lazily-created pipelines are compiled once per unique state/format
 * combination. One cache per context; no global registry.
 *
 * Pipelines are keyed by nested maps on the source strings themselves rather
 * than content hashes: lookups never walk the (potentially large) WGSL source
 * and keys cannot collide.
 */
export class PipelineCache {
  #device: GPUDevice;
  #modules = new Map<string, GPUShaderModule>();
  /** Vertex source → fragment source → serialized state */
  #renderPipelines = new Map<
    string,
    Map<string, Map<string, GPURenderPipeline>>
  >();
  /** Compute source → serialized state */
  #computePipelines = new Map<string, Map<string, GPUComputePipeline>>();
  #failed = new WeakSet<GPURenderPipeline | GPUComputePipeline>();

  constructor(device: GPUDevice) {
    this.#device = device;
  }

  /**
   * Time already charged to the nested caches a pipeline creation goes through,
   * subtracted from the pipeline's own `ms` so the counters stay disjoint and
   * additive. The lifetime counters are monotonic, so the delta across a
   * factory call is exactly what that call triggered.
   */
  #nestedTime(): number {
    const { lifetime } = debugCounters(this.#device);
    return lifetime.shaderReflections.ms + lifetime.shaderModules.ms;
  }

  get size(): number {
    let size = 0;
    for (const byFragment of this.#renderPipelines.values()) {
      for (const byState of byFragment.values()) size += byState.size;
    }
    for (const byState of this.#computePipelines.values()) size += byState.size;
    return size;
  }

  /**
   * Get (or create and cache) a shader module. Compilation diagnostics are
   * reported asynchronously and identically across browsers: the validation
   * error is captured in an error scope (Chrome logs uncaptured ones with its
   * own formatting, Safari not at all) and getCompilationInfo() messages are
   * formatted with the offending source line and a caret. Reporting never
   * blocks module creation, and runs once per unique source.
   */
  getShaderModule(code: string, label?: string): GPUShaderModule {
    let created = false;
    let start = 0;
    const module = this.#modules.getOrInsertComputed(code, (code) => {
      created = true;
      start = performance.now();
      debugLog(
        this.#device,
        "resources",
        () =>
          `createShaderModule ${label ? `"${label}" ` : ""}(${code.length} chars)`,
      );
      this.#device.pushErrorScope("validation");
      const module = this.#device.createShaderModule({
        code,
        ...(label && { label }),
      });
      // popErrorScope must be called synchronously — scopes are a LIFO stack
      // and a deferred call would pop an unrelated (eg. pipeline) scope.
      void reportShaderDiagnostics(
        code,
        module,
        this.#device.popErrorScope(),
        label,
      );
      return module;
    });
    recordCacheAccess(
      this.#device,
      "shaderModules",
      created ? performance.now() - start : null,
    );
    return module;
  }

  /**
   * Browsers differ in when WGSL errors surface: Chrome reports at module
   * creation, WebKit defers compilation so the validation error can fire at
   * pipeline creation instead — capture that too, or Safari stays silent.
   */
  #captureValidation<T extends GPURenderPipeline | GPUComputePipeline>(
    label: string,
    create: () => T,
    stages: StageModule[],
  ): T {
    this.#device.pushErrorScope("validation");
    const pipeline = create();
    void reportPipelineError(this.#device.popErrorScope(), label, stages).then(
      (failed) => {
        if (failed) this.#failed.add(pipeline);
      },
    );
    return pipeline;
  }

  /**
   * Whether a pipeline from this cache failed validation. Known asynchronously:
   * false until its error scope resolves, so the first submits using it are
   * already encoded.
   */
  hasFailed(pipeline: GPURenderPipeline | GPUComputePipeline): boolean {
    return this.#failed.has(pipeline);
  }

  getRenderPipeline(source: RenderPipelineSource): GPURenderPipeline {
    // Normalized before keying, so `shader` and an identical vertex/fragment
    // pair land in the same cache bucket.
    const { vertex, fragment } = resolveShaders(source);
    const byFragment = this.#renderPipelines.getOrInsertComputed(
      vertex,
      () => new Map(),
    );
    const byState = byFragment.getOrInsertComputed(
      fragment ?? "",
      () => new Map(),
    );
    const state = JSON.stringify([
      source.buffers,
      source.targets,
      source.primitive,
      source.depthStencil,
      source.multisample,
      source.vertexEntryPoint,
      source.fragmentEntryPoint,
      source.constants,
      source.layoutKey ?? "auto",
    ]);

    let created = false;
    let start = 0;
    let nestedTime = 0;
    const pipeline = byState.getOrInsertComputed(state, () => {
      created = true;
      start = performance.now();
      nestedTime = this.#nestedTime();
      const label = source.label ?? "pex-gpu pipeline";
      debugLog(
        this.#device,
        "resources",
        () => `createRenderPipeline "${label}" ${state}`,
      );
      const vertexModule = this.getShaderModule(vertex, source.label);
      // A combined `shader` declares both stages from one source: looking it
      // up twice would book a guaranteed hit against every real miss.
      let fragmentModule: GPUShaderModule | undefined;
      if (fragment) {
        fragmentModule =
          fragment === vertex
            ? vertexModule
            : this.getShaderModule(fragment, source.label);
      }

      // parseWGSL, not reflectWGSL: these are entry point and override
      // lookups, not a reflection request worth counting. The commands layer
      // has already reflected both sources by the time it gets here, so
      // recording them only buried the real miss under internal hits.
      const vertexReflection = parseWGSL(vertex);
      let fragmentReflection: ShaderReflection | undefined;
      if (fragment) {
        fragmentReflection =
          fragment === vertex ? vertexReflection : parseWGSL(fragment);
      }
      const vertexEntryPoint =
        source.vertexEntryPoint ?? vertexReflection.entryPoints.vertex;
      const fragmentEntryPoint =
        source.fragmentEntryPoint ?? fragmentReflection?.entryPoints.fragment;

      const vertexConstants = scopeConstants(
        source.constants,
        vertexEntryPoint
          ? vertexReflection.overrides.get(vertexEntryPoint)
          : undefined,
      );
      const fragmentConstants =
        fragmentReflection && fragmentEntryPoint
          ? scopeConstants(
              source.constants,
              fragmentReflection.overrides.get(fragmentEntryPoint),
            )
          : undefined;

      const stages: StageModule[] = [
        { stage: "vertex", code: vertex, module: vertexModule },
      ];
      if (fragmentModule && fragment) {
        stages.push({
          stage: "fragment",
          code: fragment,
          module: fragmentModule,
        });
      }

      return this.#captureValidation(
        label,
        () =>
          this.#device.createRenderPipeline({
            label,
            layout: source.layout ?? "auto",
            vertex: {
              module: vertexModule,
              ...(vertexEntryPoint && { entryPoint: vertexEntryPoint }),
              ...(vertexConstants && { constants: vertexConstants }),
              buffers: source.buffers ?? [],
            },
            ...(fragmentModule && {
              fragment: {
                module: fragmentModule,
                ...(fragmentEntryPoint && { entryPoint: fragmentEntryPoint }),
                ...(fragmentConstants && { constants: fragmentConstants }),
                targets: source.targets ?? [],
              },
            }),
            ...(source.primitive && { primitive: source.primitive }),
            ...(source.depthStencil && { depthStencil: source.depthStencil }),
            ...(source.multisample && { multisample: source.multisample }),
          }),
        stages,
      );
    });
    recordCacheAccess(
      this.#device,
      "renderPipelines",
      created
        ? performance.now() - start - (this.#nestedTime() - nestedTime)
        : null,
    );
    return pipeline;
  }

  getComputePipeline(source: ComputePipelineSource): GPUComputePipeline {
    const byState = this.#computePipelines.getOrInsertComputed(
      source.compute,
      () => new Map(),
    );
    const state = JSON.stringify([
      source.entryPoint,
      source.constants,
      source.layoutKey ?? "auto",
    ]);

    let created = false;
    let start = 0;
    let nestedTime = 0;
    const pipeline = byState.getOrInsertComputed(state, () => {
      created = true;
      start = performance.now();
      nestedTime = this.#nestedTime();
      const label = source.label ?? "pex-gpu compute pipeline";
      debugLog(
        this.#device,
        "resources",
        () => `createComputePipeline "${label}" ${state}`,
      );
      const reflection = parseWGSL(source.compute);
      const entryPoint = source.entryPoint ?? reflection.entryPoints.compute;
      const module = this.getShaderModule(source.compute, source.label);
      const constants = scopeConstants(
        source.constants,
        entryPoint ? reflection.overrides.get(entryPoint) : undefined,
      );
      return this.#captureValidation(
        label,
        () =>
          this.#device.createComputePipeline({
            label,
            layout: source.layout ?? "auto",
            compute: {
              module,
              ...(entryPoint && { entryPoint }),
              ...(constants && { constants }),
            },
          }),
        [{ stage: "compute", code: source.compute, module }],
      );
    });
    recordCacheAccess(
      this.#device,
      "computePipelines",
      created
        ? performance.now() - start - (this.#nestedTime() - nestedTime)
        : null,
    );
    return pipeline;
  }

  clear(): void {
    this.#modules.clear();
    this.#renderPipelines.clear();
    this.#computePipelines.clear();
  }
}