/* A small WebGL renderer for the X2's URDF model.

   Deliberately hand-rolled rather than pulled from a 3D library. What this
   needs is narrow - one shader, flat shading, a rigid-body tree of 41 links -
   and the alternative is several hundred kilobytes fetched over the robot's
   own Wi-Fi, from a CDN the robot cannot reach.

   Geometry arrives as uint16 positions quantised inside each mesh's bounding
   box (see build_model.py); the shader expands them back to metres, so the
   4x saving over float32 costs two extra instructions per vertex.

   Normals are computed per-face in the fragment shader from screen-space
   derivatives. That means the vertex buffer carries positions and nothing
   else - no normal buffer, no smoothing groups, no index-order assumptions -
   and flat shading is the honest look for a machined-metal robot anyway.
*/

const VERTEX_SHADER = `
attribute vec3 a_quantised;
uniform mat4 u_viewProjection;
uniform mat4 u_model;
uniform vec3 u_lo;
uniform vec3 u_span;
varying vec3 v_world;
void main() {
  vec3 local = u_lo + (a_quantised / 65535.0) * u_span;
  vec4 world = u_model * vec4(local, 1.0);
  v_world = world.xyz;
  gl_Position = u_viewProjection * world;
}`;

const FRAGMENT_SHADER = `
precision mediump float;
varying vec3 v_world;
uniform vec3 u_colour;
uniform float u_alpha;
void main() {
  // Face normal straight from the derivative of world position across the
  // triangle. No normal attribute needed, and it cannot disagree with the
  // geometry the way a stale baked normal can.
  vec3 normal = normalize(cross(dFdx(v_world), dFdy(v_world)));
  vec3 keyDir = normalize(vec3(0.45, 0.7, 0.85));
  vec3 fillDir = normalize(vec3(-0.6, -0.3, 0.4));
  float key = max(dot(normal, keyDir), 0.0);
  float fill = max(dot(normal, fillDir), 0.0) * 0.35;
  // A little rim light so the silhouette stays readable against a dark panel.
  float rim = pow(1.0 - abs(normal.z), 2.0) * 0.18;
  vec3 shaded = u_colour * (0.30 + 0.72 * key + fill) + rim;
  gl_FragColor = vec4(shaded, u_alpha);
}`;

/* -- Small matrix helpers (column-major, as WebGL expects) ----------------- */

export function identity() {
  return new Float32Array([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]);
}

export function multiply(a, b) {
  const out = new Float32Array(16);
  for (let col = 0; col < 4; col += 1) {
    for (let row = 0; row < 4; row += 1) {
      out[col * 4 + row] =
        a[row] * b[col * 4] +
        a[4 + row] * b[col * 4 + 1] +
        a[8 + row] * b[col * 4 + 2] +
        a[12 + row] * b[col * 4 + 3];
    }
  }
  return out;
}

/** URDF fixed-axis roll-pitch-yaw: R = Rz(yaw) * Ry(pitch) * Rx(roll). */
export function fromOrigin(xyz, rpy) {
  const [x, y, z] = xyz;
  const [roll, pitch, yaw] = rpy;
  const cr = Math.cos(roll), sr = Math.sin(roll);
  const cp = Math.cos(pitch), sp = Math.sin(pitch);
  const cy = Math.cos(yaw), sy = Math.sin(yaw);

  return new Float32Array([
    cy * cp, sy * cp, -sp, 0,
    cy * sp * sr - sy * cr, sy * sp * sr + cy * cr, cp * sr, 0,
    cy * sp * cr + sy * sr, sy * sp * cr - cy * sr, cp * cr, 0,
    x, y, z, 1,
  ]);
}

/** Rotation of `angle` about an arbitrary unit axis (Rodrigues). */
export function fromAxisAngle(axis, angle) {
  let [x, y, z] = axis;
  const length = Math.hypot(x, y, z) || 1;
  x /= length; y /= length; z /= length;
  const c = Math.cos(angle), s = Math.sin(angle), t = 1 - c;

  return new Float32Array([
    t * x * x + c, t * x * y + s * z, t * x * z - s * y, 0,
    t * x * y - s * z, t * y * y + c, t * y * z + s * x, 0,
    t * x * z + s * y, t * y * z - s * x, t * z * z + c, 0,
    0, 0, 0, 1,
  ]);
}

export function perspective(fovY, aspect, near, far) {
  const f = 1 / Math.tan(fovY / 2);
  return new Float32Array([
    f / aspect, 0, 0, 0,
    0, f, 0, 0,
    0, 0, (far + near) / (near - far), -1,
    0, 0, (2 * far * near) / (near - far), 0,
  ]);
}

export function lookAt(eye, target, up) {
  const zx = eye[0] - target[0], zy = eye[1] - target[1], zz = eye[2] - target[2];
  const zl = Math.hypot(zx, zy, zz) || 1;
  const z = [zx / zl, zy / zl, zz / zl];

  const xx = up[1] * z[2] - up[2] * z[1];
  const xy = up[2] * z[0] - up[0] * z[2];
  const xz = up[0] * z[1] - up[1] * z[0];
  const xl = Math.hypot(xx, xy, xz) || 1;
  const x = [xx / xl, xy / xl, xz / xl];

  const y = [
    z[1] * x[2] - z[2] * x[1],
    z[2] * x[0] - z[0] * x[2],
    z[0] * x[1] - z[1] * x[0],
  ];

  return new Float32Array([
    x[0], y[0], z[0], 0,
    x[1], y[1], z[1], 0,
    x[2], y[2], z[2], 0,
    -(x[0] * eye[0] + x[1] * eye[1] + x[2] * eye[2]),
    -(y[0] * eye[0] + y[1] * eye[1] + y[2] * eye[2]),
    -(z[0] * eye[0] + z[1] * eye[1] + z[2] * eye[2]),
    1,
  ]);
}

/* -- Renderer -------------------------------------------------------------- */

export class RobotModel {
  constructor(canvas) {
    this.canvas = canvas;
    this.gl = canvas.getContext('webgl', {
      antialias: true, alpha: false, depth: true, preserveDrawingBuffer: false,
    });
    if (!this.gl) throw new Error('This browser has no WebGL support.');

    // Flat shading needs dFdx/dFdy, which is an extension in WebGL 1. Without
    // it the shader still links on most drivers but every face comes out
    // uniformly lit, so say so rather than render something misleading.
    this.derivatives = this.gl.getExtension('OES_standard_derivatives');

    this.model = null;
    this.parts = [];
    this.jointAngles = new Map();
    this.jointsByChild = new Map();
    this.linkParent = new Map();
    this.colour = [0.62, 0.66, 0.72];
    this.accent = [0.22, 0.53, 0.90];
    this.highlight = new Set();
    // link name -> [r, g, b]. Anything not listed falls back to this.colour, so
    // a link whose joint reports nothing stays neutral grey rather than
    // rendering as "zero load", which would read as a measurement.
    this.linkColours = new Map();

    this._buildProgram();
  }

  _buildProgram() {
    const gl = this.gl;
    const prefix = this.derivatives ? '#extension GL_OES_standard_derivatives : enable\n' : '';

    const compile = (type, source) => {
      const shader = gl.createShader(type);
      gl.shaderSource(shader, source);
      gl.compileShader(shader);
      if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
        throw new Error(`Shader failed to compile: ${gl.getShaderInfoLog(shader)}`);
      }
      return shader;
    };

    const program = gl.createProgram();
    gl.attachShader(program, compile(gl.VERTEX_SHADER, VERTEX_SHADER));
    gl.attachShader(program, compile(gl.FRAGMENT_SHADER, prefix + FRAGMENT_SHADER));
    gl.linkProgram(program);
    if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
      throw new Error(`Shader failed to link: ${gl.getProgramInfoLog(program)}`);
    }

    this.program = program;
    this.attribs = { quantised: gl.getAttribLocation(program, 'a_quantised') };
    this.uniforms = {
      viewProjection: gl.getUniformLocation(program, 'u_viewProjection'),
      model: gl.getUniformLocation(program, 'u_model'),
      lo: gl.getUniformLocation(program, 'u_lo'),
      span: gl.getUniformLocation(program, 'u_span'),
      colour: gl.getUniformLocation(program, 'u_colour'),
      alpha: gl.getUniformLocation(program, 'u_alpha'),
    };
  }

  /** Upload the baked model. `geometry` is the raw model.bin ArrayBuffer. */
  load(model, geometry) {
    const gl = this.gl;
    this.model = model;
    this.parts = [];

    for (const joint of model.joints) {
      this.jointsByChild.set(joint.child, joint);
      this.linkParent.set(joint.child, joint.parent);
      if (joint.type === 'revolute') this.jointAngles.set(joint.name, 0);
    }

    for (const [linkName, link] of Object.entries(model.links)) {
      for (const mesh of link.meshes) {
        const vertexBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
        gl.bufferData(gl.ARRAY_BUFFER,
          new Uint16Array(geometry, mesh.vertex_offset, mesh.vertex_count * 3),
          gl.STATIC_DRAW);

        const indexBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, indexBuffer);
        const indices = mesh.index_bits === 16
          ? new Uint16Array(geometry, mesh.index_offset, mesh.index_count)
          : new Uint32Array(geometry, mesh.index_offset, mesh.index_count);
        gl.bufferData(gl.ELEMENT_ARRAY_BUFFER, indices, gl.STATIC_DRAW);

        this.parts.push({
          link: linkName,
          vertexBuffer,
          indexBuffer,
          count: mesh.index_count,
          type: mesh.index_bits === 16 ? gl.UNSIGNED_SHORT : gl.UNSIGNED_INT,
          lo: mesh.lo,
          span: mesh.span,
          local: fromOrigin(mesh.origin.xyz, mesh.origin.rpy),
        });
      }
    }

    // 32-bit indices are an extension in WebGL 1. Every mesh here is under
    // 65535 vertices so it should never come up, but a silent wrong-geometry
    // render is worse than a clear failure.
    if (this.parts.some((part) => part.type === gl.UNSIGNED_INT)
        && !gl.getExtension('OES_element_index_uint')) {
      throw new Error('This browser cannot draw 32-bit indices (OES_element_index_uint).');
    }
  }

  setJoint(name, radians) {
    if (this.jointAngles.has(name)) this.jointAngles.set(name, radians);
  }

  setJoints(values) {
    for (const [name, radians] of Object.entries(values)) this.setJoint(name, radians);
  }

  /** World transform for a link, walking up to the root through its joints. */
  worldTransform(linkName) {
    const chain = [];
    let current = linkName;
    // Guarded against a malformed tree: a cycle would otherwise hang the frame.
    for (let depth = 0; current && depth < 64; depth += 1) {
      const joint = this.jointsByChild.get(current);
      if (!joint) break;
      chain.push(joint);
      current = joint.parent;
    }

    let matrix = identity();
    for (const joint of chain) {
      let local = fromOrigin(joint.origin.xyz, joint.origin.rpy);
      if (joint.type === 'revolute') {
        const angle = this.jointAngles.get(joint.name) || 0;
        local = multiply(local, fromAxisAngle(joint.axis, angle));
      }
      matrix = multiply(local, matrix);
    }
    return matrix;
  }

  resize() {
    const gl = this.gl;
    // Cap the backing store at 2x CSS pixels: past that the extra fragments
    // cost real milliseconds on a phone and buy nothing visible.
    const ratio = Math.min(window.devicePixelRatio || 1, 2);
    const width = Math.round(this.canvas.clientWidth * ratio);
    const height = Math.round(this.canvas.clientHeight * ratio);
    if (width && height && (this.canvas.width !== width || this.canvas.height !== height)) {
      this.canvas.width = width;
      this.canvas.height = height;
    }
    gl.viewport(0, 0, this.canvas.width, this.canvas.height);
  }

  render(camera, background = [0.09, 0.10, 0.12]) {
    const gl = this.gl;
    this.resize();

    gl.clearColor(background[0], background[1], background[2], 1);
    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);
    if (!this.model) return;

    gl.enable(gl.DEPTH_TEST);
    // Backface culling stays off: vertex clustering does not preserve winding
    // reliably, so culling by winding punches holes in the model. Depth testing
    // alone is correct here; the cost is shading some faces we then overwrite.
    gl.disable(gl.CULL_FACE);

    const aspect = this.canvas.width / Math.max(1, this.canvas.height);
    const eye = [
      camera.target[0] + camera.distance * Math.cos(camera.pitch) * Math.cos(camera.yaw),
      camera.target[1] + camera.distance * Math.cos(camera.pitch) * Math.sin(camera.yaw),
      camera.target[2] + camera.distance * Math.sin(camera.pitch),
    ];
    const viewProjection = multiply(
      perspective(camera.fov ?? 0.8, aspect, 0.05, 60),
      // Z is up in URDF, so the camera's up vector is Z, not Y.
      lookAt(eye, camera.target, [0, 0, 1]),
    );

    gl.useProgram(this.program);
    gl.uniformMatrix4fv(this.uniforms.viewProjection, false, viewProjection);
    gl.enableVertexAttribArray(this.attribs.quantised);

    const worldCache = new Map();

    for (const part of this.parts) {
      let world = worldCache.get(part.link);
      if (!world) {
        world = this.worldTransform(part.link);
        worldCache.set(part.link, world);
      }

      gl.uniformMatrix4fv(this.uniforms.model, false, multiply(world, part.local));
      gl.uniform3fv(this.uniforms.lo, part.lo);
      gl.uniform3fv(this.uniforms.span, part.span);

      const tinted = this.linkColours.get(part.link);
      const colour = tinted || (this.highlight.has(part.link) ? this.accent : this.colour);
      gl.uniform3fv(this.uniforms.colour, colour);
      gl.uniform1f(this.uniforms.alpha, 1.0);

      gl.bindBuffer(gl.ARRAY_BUFFER, part.vertexBuffer);
      // normalized=false: the shader divides by 65535 itself, because
      // normalising here would also apply to the index-style values and makes
      // the dequantisation harder to follow.
      gl.vertexAttribPointer(this.attribs.quantised, 3, gl.UNSIGNED_SHORT, false, 0, 0);
      gl.bindBuffer(gl.ELEMENT_ARRAY_BUFFER, part.indexBuffer);
      gl.drawElements(gl.TRIANGLES, part.count, part.type, 0);
    }
  }

  dispose() {
    const gl = this.gl;
    for (const part of this.parts) {
      gl.deleteBuffer(part.vertexBuffer);
      gl.deleteBuffer(part.indexBuffer);
    }
    this.parts = [];
    gl.deleteProgram(this.program);
    // Free the drawing buffer immediately rather than waiting for GC. Browsers
    // cap the number of live WebGL contexts (often 16), and navigating between
    // tabs repeatedly would otherwise start losing them.
    gl.getExtension('WEBGL_lose_context')?.loseContext();
  }
}