import { forwardRef, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react";
import {
  colormapTable,
  getColormap,
  holePosition,
  type FieldFrames,
  type PlateResult,
  type ProblemSpec,
  type Range,
} from "@plate-dyn/core";

export interface Plate3DViewProps {
  result: PlateResult;
  spec: ProblemSpec;
  frameIndex: number;
  /** Frames that displace the surface along the normal (the flexural field, or a mode shape) */
  frames?: FieldFrames;
  /** Frames that colour the surface (any per-node field) and their colour range */
  colorFrames?: FieldFrames;
  colorRange?: Range;
  colormap?: string;
  /** Exaggeration of the displacement */
  scale: number;
  width?: number;
  height?: number;
}

type Vec3 = [number, number, number];

const VS = `
attribute vec3 aPos;
attribute vec3 aNormal;
attribute vec3 aColor;
uniform mat4 uProj;
uniform mat4 uView;
varying vec3 vColor;
varying vec3 vNormal;
void main() {
  gl_Position = uProj * uView * vec4(aPos, 1.0);
  vColor = aColor;
  vNormal = (uView * vec4(aNormal, 0.0)).xyz;
}
`;
const FS = `
precision mediump float;
varying vec3 vColor;
varying vec3 vNormal;
uniform float uFlat;
void main() {
  vec3 n = normalize(vNormal);
  float diff = abs(dot(n, normalize(vec3(0.3, 0.6, 1.0))));
  float shade = mix(0.55 + 0.45 * diff, 1.0, uFlat);
  gl_FragColor = vec4(vColor * shade, 1.0);
}
`;

function perspective(fovy: number, aspect: number, near: number, far: number): Float32Array {
  const f = 1 / Math.tan(fovy / 2);
  const out = new Float32Array(16);
  out[0] = f / aspect;
  out[5] = f;
  out[10] = (far + near) / (near - far);
  out[11] = -1;
  out[14] = (2 * far * near) / (near - far);
  return out;
}

function lookAt(eye: Vec3, center: Vec3, up: Vec3): Float32Array {
  const z = norm(sub(eye, center));
  const x = norm(cross(up, z));
  const y = cross(z, x);
  const out = new Float32Array(16);
  out[0] = x[0];
  out[4] = x[1];
  out[8] = x[2];
  out[1] = y[0];
  out[5] = y[1];
  out[9] = y[2];
  out[2] = z[0];
  out[6] = z[1];
  out[10] = z[2];
  out[12] = -dot(x, eye);
  out[13] = -dot(y, eye);
  out[14] = -dot(z, eye);
  out[15] = 1;
  return out;
}
const sub = (a: Vec3, b: Vec3): Vec3 => [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
const add = (a: Vec3, b: Vec3): Vec3 => [a[0] + b[0], a[1] + b[1], a[2] + b[2]];
const dot = (a: Vec3, b: Vec3) => a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
const cross = (a: Vec3, b: Vec3): Vec3 => [
  a[1] * b[2] - a[2] * b[1],
  a[2] * b[0] - a[0] * b[2],
  a[0] * b[1] - a[1] * b[0],
];
const norm = (a: Vec3): Vec3 => {
  const l = Math.hypot(a[0], a[1], a[2]) || 1;
  return [a[0] / l, a[1] / l, a[2] / l];
};

/** What the movie recorder needs of the scene: its canvas, and a frame drawn on demand. */
export interface Plate3DHandle {
  canvas: () => HTMLCanvasElement | null;
  /** Draw frame k; `size` overrides the backing store, for recording at a larger size */
  drawFrame: (k: number, size?: { width: number; height: number }) => void;
}

/**
 * The plate as it hangs: the thread rising from the hole near the top edge,
 * the hammer's tip at the strike point on the near face, the surface
 * displaced along the normal by the chosen frames (exaggerated by `scale`)
 * and coloured by the chosen quantity. The plate's x runs across the screen
 * and its y up it. Drag to orbit, wheel to zoom.
 */
export const Plate3DView = forwardRef<Plate3DHandle, Plate3DViewProps>(
  function Plate3DView(props, ref) {
    const {
      result,
      spec,
      frameIndex,
      frames,
      colorFrames,
      colorRange,
      colormap = "bwr",
      scale,
      width = 460,
      height = 340,
    } = props;
    const canvasRef = useRef<HTMLCanvasElement>(null);
    const glRef = useRef<{
      gl: WebGLRenderingContext;
      prog: WebGLProgram;
      buffers: { pos: WebGLBuffer; nrm: WebGLBuffer; col: WebGLBuffer; idx: WebGLBuffer };
      lines: { pos: WebGLBuffer; col: WebGLBuffer; nrm: WebGLBuffer };
      loc: Record<string, number | WebGLUniformLocation | null>;
    } | null>(null);
    const [cam, setCam] = useState({ az: 0.5, el: 0.3, zoom: 1 });
    const drag = useRef<{ x: number; y: number } | null>(null);

    const nx = result.grid.x.length;
    const ny = result.grid.y.length;
    const W = spec.plate.width;
    const H = spec.plate.height;
    const size = Math.max(W, H);

    const mesh = useMemo(() => {
      const idx: number[] = [];
      for (let j = 0; j < nx - 1; j++)
        for (let i = 0; i < ny - 1; i++) {
          const a = i + j * ny;
          const b = i + 1 + j * ny;
          const c = i + (j + 1) * ny;
          const d = i + 1 + (j + 1) * ny;
          idx.push(a, b, c, b, d, c);
        }
      return { index: new Uint16Array(idx) };
    }, [nx, ny]);

    useEffect(() => {
      const canvas = canvasRef.current;
      if (!canvas) return;
      const gl = canvas.getContext("webgl", { antialias: true });
      if (!gl) return;
      const compile = (type: number, src: string) => {
        const sh = gl.createShader(type)!;
        gl.shaderSource(sh, src);
        gl.compileShader(sh);
        return sh;
      };
      const prog = gl.createProgram()!;
      gl.attachShader(prog, compile(gl.VERTEX_SHADER, VS));
      gl.attachShader(prog, compile(gl.FRAGMENT_SHADER, FS));
      gl.linkProgram(prog);
      const mk = () => gl.createBuffer()!;
      glRef.current = {
        gl,
        prog,
        buffers: { pos: mk(), nrm: mk(), col: mk(), idx: mk() },
        lines: { pos: mk(), col: mk(), nrm: mk() },
        loc: {
          aPos: gl.getAttribLocation(prog, "aPos"),
          aNormal: gl.getAttribLocation(prog, "aNormal"),
          aColor: gl.getAttribLocation(prog, "aColor"),
          uProj: gl.getUniformLocation(prog, "uProj"),
          uView: gl.getUniformLocation(prog, "uView"),
          uFlat: gl.getUniformLocation(prog, "uFlat"),
        },
      };
      return () => {
        glRef.current = null;
      };
    }, []);

    const draw = (kIndex: number, out?: { width: number; height: number }) => {
      const r = glRef.current;
      const canvas = canvasRef.current;
      if (!r || !canvas) return;
      const { gl, prog, buffers, lines, loc } = r;
      const vw = out?.width ?? width;
      const vh = out?.height ?? height;
      const dpr = out ? 1 : Math.min(2, window.devicePixelRatio || 1);
      canvas.width = Math.round(vw * dpr);
      canvas.height = Math.round(vh * dpr);
      gl.viewport(0, 0, canvas.width, canvas.height);

      // ---- node positions: plate x across, y up, the normal towards the viewer
      const N = nx * ny;
      const kf = frames ? Math.min(kIndex, frames.nf - 1) : -1;
      const pos = new Float32Array(3 * N);
      const labPos: Vec3[] = new Array(N);
      for (let j = 0; j < nx; j++)
        for (let i = 0; i < ny; i++) {
          const n = i + j * ny;
          const f = frames && kf >= 0 ? frames.data[n + kf * N] : 0;
          const lp: Vec3 = [
            result.grid.x[j],
            result.grid.y[i],
            scale * (Number.isFinite(f) ? f : 0),
          ];
          labPos[n] = lp;
          pos[3 * n] = lp[0];
          pos[3 * n + 1] = lp[1];
          pos[3 * n + 2] = lp[2];
        }
      const nrm = new Float32Array(3 * N);
      for (let j = 0; j < nx; j++)
        for (let i = 0; i < ny; i++) {
          const n = i + j * ny;
          const jl = Math.max(0, j - 1);
          const jr = Math.min(nx - 1, j + 1);
          const il = Math.max(0, i - 1);
          const ir = Math.min(ny - 1, i + 1);
          const tx = sub(labPos[i + jr * ny], labPos[i + jl * ny]);
          const ty = sub(labPos[ir + j * ny], labPos[il + j * ny]);
          const nn = norm(cross(tx, ty));
          nrm[3 * n] = nn[0];
          nrm[3 * n + 1] = nn[1];
          nrm[3 * n + 2] = nn[2];
        }
      const col = new Float32Array(3 * N);
      const table = colormapTable(getColormap(colormap));
      const range = colorRange ?? { min: -1, max: 1 };
      const span = range.max - range.min || 1;
      const kc = colorFrames ? Math.min(kIndex, colorFrames.nf - 1) : -1;
      for (let n = 0; n < N; n++) {
        let v = colorFrames && kc >= 0 ? colorFrames.data[n + kc * N] : NaN;
        if (!Number.isFinite(v)) v = 0.5 * (range.min + range.max);
        let t = (v - range.min) / span;
        t = t < 0 ? 0 : t > 1 ? 1 : t;
        const k = Math.round(t * 255) * 3;
        col[3 * n] = table[k] / 255;
        col[3 * n + 1] = table[k + 1] / 255;
        col[3 * n + 2] = table[k + 2] / 255;
      }

      // ---- the thread, the hammer, the outline
      const L: number[] = [];
      const LC: number[] = [];
      const line = (a: Vec3, b: Vec3, rgb: Vec3) => {
        L.push(...a, ...b);
        LC.push(...rgb, ...rgb);
      };
      const dark: Vec3 = [0.15, 0.15, 0.18];
      const red: Vec3 = [0.82, 0.29, 0.36];
      const pale: Vec3 = [0.75, 0.75, 0.78];
      const grey: Vec3 = [0.4, 0.4, 0.45];
      const circle = (centre: Vec3, r: number, rgb: Vec3, axis: "z" | "y" = "z") => {
        const n = 24;
        for (let k = 0; k < n; k++) {
          const a0 = (2 * Math.PI * k) / n;
          const a1 = (2 * Math.PI * (k + 1)) / n;
          const p = (a: number): Vec3 =>
            axis === "z"
              ? [centre[0] + r * Math.cos(a), centre[1] + r * Math.sin(a), centre[2]]
              : [centre[0] + r * Math.cos(a), centre[1], centre[2] + r * Math.sin(a)];
          line(p(a0), p(a1), rgb);
        }
      };
      const hole = holePosition(spec);
      const hookUp = Math.min(spec.thread.length, 0.7 * size);
      line([hole.x, hole.y, 0], [hole.x, hole.y + hookUp, 0], dark);
      line(
        [hole.x - 0.03 * size, hole.y + hookUp, 0],
        [hole.x + 0.03 * size, hole.y + hookUp, 0],
        dark,
      );
      circle([hole.x, hole.y, 0], spec.hole.diameter / 2, dark);
      // the hammer: its tip resting on the near face, its shank along the normal
      const R = spec.hammer.tip_radius;
      const s: Vec3 = [spec.strike.x, spec.strike.y, 0];
      circle([s[0], s[1], R], R, red);
      circle([s[0], s[1], R], R, red, "y");
      line([s[0], s[1], 2 * R], [s[0], s[1], 2 * R + 0.35 * size], red);
      // the probe
      line(
        [spec.probe.x - 0.02 * size, spec.probe.y, 0],
        [spec.probe.x + 0.02 * size, spec.probe.y, 0],
        grey,
      );
      line(
        [spec.probe.x, spec.probe.y - 0.02 * size, 0],
        [spec.probe.x, spec.probe.y + 0.02 * size, 0],
        grey,
      );
      // outline of the undeformed plate, faint
      line([-W / 2, -H / 2, 0], [W / 2, -H / 2, 0], pale);
      line([W / 2, -H / 2, 0], [W / 2, H / 2, 0], pale);
      line([W / 2, H / 2, 0], [-W / 2, H / 2, 0], pale);
      line([-W / 2, H / 2, 0], [-W / 2, -H / 2, 0], pale);

      // ---- camera
      const centre: Vec3 = [0, 0.08 * size, 0];
      const dist = (2.4 * size) / cam.zoom;
      const eye: Vec3 = add(centre, [
        dist * Math.cos(cam.el) * Math.sin(cam.az),
        dist * Math.sin(cam.el),
        dist * Math.cos(cam.el) * Math.cos(cam.az),
      ]);
      const view = lookAt(eye, centre, [0, 1, 0]);
      const proj = perspective(0.6, vw / vh, 0.01 * size, 20 * size);

      gl.clearColor(0.97, 0.97, 0.975, 1);
      gl.enable(gl.DEPTH_TEST);
      gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);
      gl.useProgram(prog);
      gl.uniformMatrix4fv(loc.uProj as WebGLUniformLocation, false, proj);
      gl.uniformMatrix4fv(loc.uView as WebGLUniformLocation, false, view);
      const attr = (buf: WebGLBuffer, data: Float32Array, a: number) => {
        gl.bindBuffer(gl.ARRAY_BUFFER, buf);
        gl.bufferData(gl.ARRAY_BUFFER, data, gl.DYNAMIC_DRAW);
        gl.enableVertexAttribArray(a);
        gl.vertexAttribPointer(a, 3, gl.FLOAT, false, 0, 0);
      };
      gl.uniform1f(loc.uFlat as WebGLUniformLocation, 0);
      attr(buffers.pos, pos, loc.aPos as number);
      attr(buffers.nrm, nrm, loc.aNormal as number);
      attr(buffers.col, col, loc.aColor as number);
      gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, buffers.idx);
      gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, mesh.index, gl.STATIC_DRAW);
      gl.drawElements(gl.TRIANGLES, mesh.index.length, gl.UNSIGNED_SHORT, 0);
      gl.uniform1f(loc.uFlat as WebGLUniformLocation, 1);
      const lp = new Float32Array(L);
      attr(lines.pos, lp, loc.aPos as number);
      attr(lines.nrm, new Float32Array(L.length), loc.aNormal as number);
      attr(lines.col, new Float32Array(LC), loc.aColor as number);
      gl.drawArrays(gl.LINES, 0, lp.length / 3);
    };

    const drawRef = useRef(draw);
    drawRef.current = draw;
    useImperativeHandle(
      ref,
      () => ({
        canvas: () => canvasRef.current,
        drawFrame: (k: number, out?: { width: number; height: number }) => drawRef.current(k, out),
      }),
      [],
    );

    useEffect(() => {
      drawRef.current(frameIndex);
    }, [
      result,
      spec,
      frameIndex,
      frames,
      colorFrames,
      colorRange,
      colormap,
      scale,
      width,
      height,
      cam,
      mesh,
      nx,
      ny,
      W,
      H,
      size,
    ]);

    return (
      <canvas
        ref={canvasRef}
        style={{ width, height, display: "block", cursor: "grab", borderRadius: 4 }}
        onPointerDown={(e) => {
          drag.current = { x: e.clientX, y: e.clientY };
          (e.target as HTMLElement).setPointerCapture(e.pointerId);
        }}
        onPointerMove={(e) => {
          if (!drag.current) return;
          const dx = e.clientX - drag.current.x;
          const dy = e.clientY - drag.current.y;
          drag.current = { x: e.clientX, y: e.clientY };
          setCam((c) => ({
            ...c,
            az: c.az - dx * 0.01,
            el: Math.max(-1.4, Math.min(1.4, c.el + dy * 0.01)),
          }));
        }}
        onPointerUp={() => {
          drag.current = null;
        }}
        onWheel={(e) => {
          e.preventDefault();
          setCam((c) => ({
            ...c,
            zoom: Math.max(0.3, Math.min(5, c.zoom * Math.exp(-e.deltaY * 0.001))),
          }));
        }}
        aria-label="the plate on its thread"
      />
    );
  },
);