import {
  LinearSRGBColorSpace,
  NearestFilter,
  NoToneMapping,
  RGBAFormat,
  Vector2,
  WebGLRenderTarget,
} from "three";
import type { Camera, Scene, WebGLRenderer } from "three";
import { EffectComposer } from "three/examples/jsm/postprocessing/EffectComposer.js";
import { RenderPass } from "three/examples/jsm/postprocessing/RenderPass.js";
import { ShaderPass } from "three/examples/jsm/postprocessing/ShaderPass.js";
import { PSXShader } from "./PSXShader.js";

/**
 * Low-resolution render target + PSX pass + nearest-neighbour upscale.
 *
 * The scene is drawn into a small target (240p by default). The final pass
 * samples it with NearestFilter and blits to the full-size canvas, which is
 * what produces the chunky upscale — doing this with CSS instead would let the
 * browser's smoothing round off exactly the edges we want to keep.
 */

export interface PSXPipelineOptions {
  /** Virtual framebuffer height. Width follows the canvas aspect ratio. */
  virtualHeight?: number;
  colorLevels?: number;
  ditherStrength?: number;
  scanlineStrength?: number;
  vignetteStrength?: number;
}

export class PSXPipeline {
  private readonly composer: EffectComposer;
  private readonly psxPass: ShaderPass;
  private readonly renderTarget: WebGLRenderTarget;
  private readonly virtualHeight: number;
  private readonly virtualSize = new Vector2(320, 240);

  constructor(
    private readonly renderer: WebGLRenderer,
    scene: Scene,
    camera: Camera,
    options: PSXPipelineOptions = {},
  ) {
    this.virtualHeight = options.virtualHeight ?? 240;

    // The PSX pass owns tone mapping and the linear→sRGB encode (see PSXShader).
    // Leaving the renderer on its defaults would either double-encode or leave
    // the canvas in linear space — both look like a black room.
    this.renderer.outputColorSpace = LinearSRGBColorSpace;
    this.renderer.toneMapping = NoToneMapping;

    this.renderTarget = new WebGLRenderTarget(320, this.virtualHeight, {
      minFilter: NearestFilter,
      magFilter: NearestFilter,
      format: RGBAFormat,
      depthBuffer: true,
      stencilBuffer: false,
      // Scene lights render in linear; keep the intermediate buffer linear so
      // the quantiser sees true radiometric values before the sRGB encode.
      colorSpace: LinearSRGBColorSpace,
    });

    this.composer = new EffectComposer(renderer, this.renderTarget);
    this.composer.addPass(new RenderPass(scene, camera));

    this.psxPass = new ShaderPass(PSXShader);
    this.psxPass.renderToScreen = true;
    // Custom fragment already wrote display-referred sRGB — do not tone-map again.
    this.psxPass.material.toneMapped = false;
    this.composer.addPass(this.psxPass);

    const uniforms = this.psxPass.uniforms;
    uniforms["uColorLevels"]!.value = options.colorLevels ?? 32;
    uniforms["uDitherStrength"]!.value = options.ditherStrength ?? 1;
    uniforms["uScanlineStrength"]!.value = options.scanlineStrength ?? 0.12;
    uniforms["uVignetteStrength"]!.value = options.vignetteStrength ?? 0.25;
  }

  /** Call on canvas resize. `width`/`height` are CSS pixels. */
  setSize(width: number, height: number): void {
    this.renderer.setSize(width, height, false);

    // Keep virtual pixels square by deriving width from the canvas aspect.
    const aspect = width / Math.max(1, height);
    const virtualWidth = Math.max(2, Math.round((this.virtualHeight * aspect) / 2) * 2);

    this.virtualSize.set(virtualWidth, this.virtualHeight);
    this.composer.setSize(virtualWidth, this.virtualHeight);
    this.psxPass.uniforms["uResolution"]!.value = this.virtualSize;
  }

  /** The grid the vertex snapper should quantise to — one snap per virtual pixel. */
  get snapResolution(): Vector2 {
    return this.virtualSize;
  }

  setCamera(camera: Camera): void {
    for (const pass of this.composer.passes) {
      if (pass instanceof RenderPass) pass.camera = camera;
    }
  }

  render(): void {
    this.composer.render();
  }

  dispose(): void {
    this.composer.dispose();
    this.renderTarget.dispose();
  }
}
