// WebGPU rendering of the displacement field: a 2D color map or a lit 3D height field.

import { type Grid, type PhysParams } from "../params.ts";
import { lookAt, multiply, perspective } from "./mat4.ts";
import shaderCode from "./shaders.wgsl?raw";

export type ViewMode = "2d" | "3d";

export class Renderer {
  mode: ViewMode = "2d";
  /** Vertical exaggeration of the 3D view. */
  exaggeration = 4;
  /** Color range |w| [m]; 0 means automatic (from the current maximum). */
  colorRange = 0;
  camera = { yaw: 0.6, pitch: 0.5, dist: 1.0 };

  private device: GPUDevice;
  private canvas: HTMLCanvasElement;
  private ctx: GPUCanvasContext;
  private format: GPUTextureFormat;
  private uniform: GPUBuffer;
  private pipe2d: GPURenderPipeline;
  private pipe3d: GPURenderPipeline;
  private bgLayout: GPUBindGroupLayout;
  private index: GPUBuffer;
  private indexCount: number;
  private depth: GPUTexture | null = null;
  private grid: Grid;
  private params: PhysParams;
  private bindGroups = new Map<GPUBuffer, GPUBindGroup>();

  constructor(device: GPUDevice, canvas: HTMLCanvasElement, grid: Grid, params: PhysParams) {
    this.device = device;
    this.canvas = canvas;
    this.grid = grid;
    this.params = params;
    this.ctx = canvas.getContext("webgpu") as GPUCanvasContext;
    this.format = navigator.gpu.getPreferredCanvasFormat();
    this.ctx.configure({ device, format: this.format, alphaMode: "opaque" });
    this.uniform = device.createBuffer({ size: 128, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    const module = device.createShaderModule({ code: shaderCode });
    this.bgLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, buffer: { type: "uniform" } },
        { binding: 1, visibility: GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, buffer: { type: "read-only-storage" } },
      ],
    });
    const layout = device.createPipelineLayout({ bindGroupLayouts: [this.bgLayout] });
    this.pipe2d = device.createRenderPipeline({
      layout,
      vertex: { module, entryPoint: "vs2d" },
      fragment: { module, entryPoint: "fs2d", targets: [{ format: this.format }] },
      primitive: { topology: "triangle-list" },
    });
    this.pipe3d = device.createRenderPipeline({
      layout,
      vertex: { module, entryPoint: "vs3d" },
      fragment: { module, entryPoint: "fs3d", targets: [{ format: this.format }] },
      primitive: { topology: "triangle-list", cullMode: "none" },
      depthStencil: { format: "depth24plus", depthWriteEnabled: true, depthCompare: "less" },
    });
    // Triangle indices of the height field.
    const { nx, ny } = grid;
    const idx = new Uint32Array((nx - 1) * (ny - 1) * 6);
    let k = 0;
    for (let j = 0; j < ny - 1; j++) {
      for (let i = 0; i < nx - 1; i++) {
        const a = j * nx + i, b = a + 1, c = a + nx, d = c + 1;
        idx[k++] = a; idx[k++] = b; idx[k++] = d;
        idx[k++] = a; idx[k++] = d; idx[k++] = c;
      }
    }
    this.indexCount = idx.length;
    this.index = device.createBuffer({ size: idx.byteLength, usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST });
    device.queue.writeBuffer(this.index, 0, idx);
    this.installControls();
  }

  private installControls(): void {
    let dragging = false, lastX = 0, lastY = 0;
    this.canvas.addEventListener("pointerdown", (e) => {
      dragging = true; lastX = e.clientX; lastY = e.clientY;
      this.canvas.setPointerCapture(e.pointerId);
    });
    this.canvas.addEventListener("pointerup", () => (dragging = false));
    this.canvas.addEventListener("pointermove", (e) => {
      if (!dragging || this.mode !== "3d") return;
      this.camera.yaw += (e.clientX - lastX) * 0.01;
      this.camera.pitch = Math.max(-1.5, Math.min(1.5, this.camera.pitch + (e.clientY - lastY) * 0.01));
      lastX = e.clientX; lastY = e.clientY;
    });
    this.canvas.addEventListener("wheel", (e) => {
      if (this.mode !== "3d") return;
      e.preventDefault();
      this.camera.dist *= Math.exp(e.deltaY * 0.001);
      this.camera.dist = Math.max(0.2, Math.min(5, this.camera.dist));
    }, { passive: false });
  }

  private resize(): void {
    const dpr = window.devicePixelRatio || 1;
    const w = Math.max(1, Math.floor(this.canvas.clientWidth * dpr));
    const h = Math.max(1, Math.floor(this.canvas.clientHeight * dpr));
    if (this.canvas.width !== w || this.canvas.height !== h || !this.depth) {
      this.canvas.width = w;
      this.canvas.height = h;
      this.depth?.destroy();
      this.depth = this.device.createTexture({ size: [w, h], format: "depth24plus", usage: GPUTextureUsage.RENDER_ATTACHMENT });
    }
  }

  /** Renders the w field held in `wBuffer` (f32 per node). `wmax` drives the automatic color range. */
  render(encoder: GPUCommandEncoder, wBuffer: GPUBuffer, wmax: number): void {
    this.resize();
    const aspect = this.canvas.width / this.canvas.height;
    const range = this.colorRange > 0 ? this.colorRange : Math.max(wmax, 1e-6);
    const u = new Float32Array(32);
    const ui = new Uint32Array(u.buffer);
    const eye = [
      this.camera.dist * Math.cos(this.camera.pitch) * Math.sin(this.camera.yaw),
      this.camera.dist * Math.sin(this.camera.pitch),
      this.camera.dist * Math.cos(this.camera.pitch) * Math.cos(this.camera.yaw),
    ];
    const mvp = multiply(perspective(0.7, aspect, 0.02, 20), lookAt(eye, [0, 0, 0], [0, 1, 0]));
    u.set(mvp, 0);
    ui[16] = this.grid.nx; ui[17] = this.grid.ny;
    u[18] = this.grid.hx; u[19] = this.grid.hy;
    u[20] = this.exaggeration; u[21] = range; u[22] = aspect;
    ui[23] = this.mode === "3d" ? 1 : 0;
    const p = this.params;
    u[24] = -p.Lx / 2 + p.holeInset; u[25] = p.Ly / 2 - p.holeInset;
    u[26] = p.Lx / 2; u[27] = -p.Ly / 2;
    u[28] = p.fingerRadius;
    u[29] = 0.4; u[30] = 0.5; u[31] = 1.0;
    this.device.queue.writeBuffer(this.uniform, 0, u);

    let bg = this.bindGroups.get(wBuffer);
    if (!bg) {
      bg = this.device.createBindGroup({
        layout: this.bgLayout,
        entries: [
          { binding: 0, resource: { buffer: this.uniform } },
          { binding: 1, resource: { buffer: wBuffer } },
        ],
      });
      this.bindGroups.set(wBuffer, bg);
    }
    const view = this.ctx.getCurrentTexture().createView();
    const clear = { r: 0.11, g: 0.12, b: 0.14, a: 1 };
    if (this.mode === "2d") {
      const pass = encoder.beginRenderPass({ colorAttachments: [{ view, clearValue: clear, loadOp: "clear", storeOp: "store" }] });
      pass.setPipeline(this.pipe2d);
      pass.setBindGroup(0, bg);
      pass.draw(3);
      pass.end();
    } else {
      const pass = encoder.beginRenderPass({
        colorAttachments: [{ view, clearValue: clear, loadOp: "clear", storeOp: "store" }],
        depthStencilAttachment: { view: this.depth!.createView(), depthClearValue: 1, depthLoadOp: "clear", depthStoreOp: "store" },
      });
      pass.setPipeline(this.pipe3d);
      pass.setBindGroup(0, bg);
      pass.setIndexBuffer(this.index, "uint32");
      pass.drawIndexed(this.indexCount);
      pass.end();
    }
  }
}