import { Vector2 } from "three";
import type { Material, Texture, WebGLProgramParametersWithUniforms } from "three";

/**
 * Per-material half of the PSX look — GDD §11.
 *
 * Two effects, both applied by patching the built-in shaders rather than
 * writing new ones, so materials keep three's lighting, fog and shadow chunks:
 *
 *  1. **Vertex snapping.** The PSX GTE had no sub-pixel precision, so vertices
 *     landed on a coarse screen grid and geometry visibly wobbled. We quantise
 *     NDC xy after projection.
 *
 *  2. **Affine texture mapping.** The console had no perspective divide per
 *     pixel, so textures warped across large triangles. GLSL ES 3.00 has no
 *     `noperspective` qualifier, so we defeat the hardware's perspective-correct
 *     interpolation arithmetically: pass `uv * w` and `w` as varyings, then
 *     divide in the fragment shader. The interpolator's own 1/w factors cancel
 *     and what survives is screen-linear — exactly the artefact we want.
 */

export interface PSXMaterialOptions {
  /** Snap grid, in virtual pixels. Lower is wobblier. */
  snapResolution?: Vector2;
  vertexSnapping?: boolean;
  affineMapping?: boolean;
}

const DEFAULT_SNAP = new Vector2(160, 120);

interface PatchedMaterial extends Material {
  userData: { psx?: { snapResolution: Vector2 } };
}

export function applyPSXMaterial<T extends Material & { map?: Texture | null }>(
  material: T,
  options: PSXMaterialOptions = {},
): T {
  const snapResolution = options.snapResolution?.clone() ?? DEFAULT_SNAP.clone();
  const vertexSnapping = options.vertexSnapping ?? true;
  // Only meaningful when there is a texture to warp; without `map` the `uv`
  // attribute may not even be declared by three's shader prefix.
  const affineMapping = (options.affineMapping ?? true) && material.map != null;

  const patched = material as unknown as PatchedMaterial;
  patched.userData.psx = { snapResolution };

  material.onBeforeCompile = (shader: WebGLProgramParametersWithUniforms) => {
    shader.uniforms["uSnapResolution"] = { value: snapResolution };

    const vertexDeclarations = [
      "uniform vec2 uSnapResolution;",
      affineMapping ? "varying vec2 vAffineUv;\nvarying float vAffineW;" : "",
    ]
      .filter(Boolean)
      .join("\n");

    const vertexBody = [
      vertexSnapping
        ? `
  {
    vec4 psxPos = gl_Position;
    psxPos.xyz /= psxPos.w;
    // Half-resolution because NDC spans -1..1 across the full width.
    vec2 grid = uSnapResolution * 0.5;
    psxPos.xy = floor(psxPos.xy * grid) / grid;
    psxPos.xyz *= psxPos.w;
    gl_Position = psxPos;
  }`
        : "",
      affineMapping
        ? `
  vAffineUv = uv * gl_Position.w;
  vAffineW = gl_Position.w;`
        : "",
    ]
      .filter(Boolean)
      .join("\n");

    shader.vertexShader = shader.vertexShader
      .replace("#include <common>", `#include <common>\n${vertexDeclarations}`)
      .replace("#include <project_vertex>", `#include <project_vertex>\n${vertexBody}`);

    if (affineMapping) {
      shader.fragmentShader = shader.fragmentShader
        .replace(
          "#include <common>",
          "#include <common>\nvarying vec2 vAffineUv;\nvarying float vAffineW;",
        )
        .replace(
          "#include <map_fragment>",
          `
  #ifdef USE_MAP
    diffuseColor *= texture2D( map, vAffineUv / vAffineW );
  #endif`,
        );
    }
  };

  // Distinct patch variants must not share a compiled program.
  material.customProgramCacheKey = () => `psx:${vertexSnapping ? 1 : 0}:${affineMapping ? 1 : 0}`;
  material.needsUpdate = true;

  return material;
}

/** Retune the snap grid at runtime — the uniform is shared by reference. */
export function setSnapResolution(material: Material, width: number, height: number): void {
  const psx = (material as PatchedMaterial).userData.psx;
  if (psx) psx.snapResolution.set(width, height);
}
