import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import type { BeadSystem, Excitation } from '../sim/types';
import { css, fmtNum } from './plot';
import { timelineBar, type Timeline } from './timeline';

// 3D view of a bead-spring system. "Setup" shows the geometry at rest with
// supports, pins, excitation and pickup marked; "Replay" animates saved
// displacement frames (exaggerated) on the shared timeline.

export interface SceneSpec {
  sys: BeadSystem;
  exc: Excitation;
  excKind: 'pluck' | 'strike' | 'hammer';
  pickup: Float64Array; // 3n weights
  pickupLabel: string;
}

export interface ReplaySpec extends SceneSpec {
  frames: Float32Array; // nFrames x 3n, matches the timeline
}

type Mode = 'setup' | 'replay';
type View = 'persp' | 'front' | 'top' | 'end';

const GLYPH = 0.12; // length of the largest excitation glyph, in units of the model size

export class BeadScene {
  private renderer: THREE.WebGLRenderer | null = null;
  private scene = new THREE.Scene();
  private camera = new THREE.PerspectiveCamera(35, 1, 0.01, 100);
  private controls: OrbitControls | null = null;
  private model = new THREE.Group();
  private wrap: HTMLDivElement;
  private legend: HTMLDivElement;
  private note: HTMLParagraphElement;
  private bar: HTMLElement;
  private exagRow: HTMLElement;
  private exagInput: HTMLInputElement;
  private tabs: HTMLElement;

  private mode: Mode = 'setup';
  private setupSpec: SceneSpec | null = null;
  private replaySpec: ReplaySpec | null = null;
  private exaggeration = 1;
  private autoExaggeration = 1;

  // geometry of the current build, in normalized scene units
  private scale = 1;
  private center = new THREE.Vector3();
  private beads: THREE.InstancedMesh | null = null;
  private pins: THREE.InstancedMesh | null = null;
  private springs: THREE.LineSegments | null = null;
  private slideRings: { mesh: THREE.Mesh; bead: number }[] = [];
  private hammerArrows: THREE.Group | null = null;
  private pickupArrow: { arrow: THREE.ArrowHelper; bead: number } | null = null;
  private spec: SceneSpec | null = null;
  private frames: Float32Array | null = null;

  constructor(root: HTMLElement, private tl: Timeline) {
    root.innerHTML = `
      <div class="card-head">
        <div class="seg" data-r="tabs"><button data-v="setup" class="on">Setup</button><button data-v="replay" disabled>Replay</button></div>
        <div class="seg" data-r="views"><button data-v="persp" class="on">3D</button><button data-v="front">Front x–y</button><button data-v="top">Top x–z</button><button data-v="end">End y–z</button></div>
      </div>
      <div class="scene-wrap" data-r="wrap"></div>
      <div class="legend scene-legend" data-r="legend"></div>
      <div class="viewer-bar" data-r="exag" hidden>
        <label class="inline">exaggeration × <input type="number" min="1" step="any" data-r="exag-in" style="width:110px"></label>
        <button class="btn-small" data-r="exag-auto">Auto</button>
      </div>
      <div data-r="bar" hidden></div>
      <p class="note" data-r="note"></p>`;
    const q = <T extends HTMLElement>(r: string) => root.querySelector(`[data-r="${r}"]`) as T;
    this.wrap = q('wrap');
    this.legend = q('legend');
    this.note = q('note');
    this.tabs = q('tabs');
    this.exagRow = q('exag');
    this.exagInput = q('exag-in');
    this.bar = timelineBar(tl);
    this.bar.hidden = true;
    q('bar').replaceWith(this.bar);

    this.tabs.addEventListener('click', (e) => {
      const b = (e.target as HTMLElement).closest('button');
      if (b && !b.disabled) this.setMode(b.dataset.v as Mode);
    });
    q('views').addEventListener('click', (e) => {
      const b = (e.target as HTMLElement).closest('button');
      if (!b) return;
      q('views').querySelectorAll('button').forEach((x) => x.classList.toggle('on', x === b));
      this.setView(b.dataset.v as View);
    });
    this.exagInput.addEventListener('change', () => {
      const v = Number(this.exagInput.value);
      if (v > 0) {
        this.exaggeration = v;
        this.update();
      }
    });
    q('exag-auto').addEventListener('click', () => {
      this.exaggeration = this.autoExaggeration;
      this.exagInput.value = fmtNum(this.exaggeration);
      this.update();
    });

    try {
      this.renderer = new THREE.WebGLRenderer({ antialias: true });
      this.renderer.setPixelRatio(window.devicePixelRatio || 1);
      this.wrap.append(this.renderer.domElement);
      this.controls = new OrbitControls(this.camera, this.renderer.domElement);
      this.controls.addEventListener('change', () => this.render());
    } catch {
      this.wrap.innerHTML = '<p class="note">WebGL is not available in this browser, so the 3D view is off.</p>';
    }
    this.scene.add(new THREE.AmbientLight(0xffffff, 1.6));
    const sun = new THREE.DirectionalLight(0xffffff, 2.2);
    sun.position.set(0.4, 1, 0.8);
    this.scene.add(sun);
    this.scene.add(this.model);
    this.setView('persp');

    tl.onChange(() => {
      if (this.mode === 'replay') this.update();
    });
    new ResizeObserver(() => this.resize()).observe(this.wrap);
    window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', () => this.rebuild());
  }

  showSetup(spec: SceneSpec) {
    this.setupSpec = spec;
    if (this.mode === 'setup') this.rebuild();
    else this.setMode('setup');
  }

  showReplay(spec: ReplaySpec) {
    this.replaySpec = spec;
    let tmax = 0;
    for (let i = 0; i < spec.frames.length; i++) if (i % 3 !== 0) tmax = Math.max(tmax, Math.abs(spec.frames[i]));
    const size = modelSize(spec.sys);
    this.autoExaggeration = tmax > 0 ? Math.max(1, (GLYPH * size) / tmax) : 1;
    this.exaggeration = this.autoExaggeration;
    this.exagInput.value = fmtNum(this.exaggeration);
    (this.tabs.querySelector('[data-v="replay"]') as HTMLButtonElement).disabled = false;
    this.setMode('replay');
  }

  private setMode(m: Mode) {
    this.mode = m;
    this.tabs.querySelectorAll<HTMLButtonElement>('button').forEach((b) => b.classList.toggle('on', b.dataset.v === m));
    this.bar.hidden = m !== 'replay';
    this.exagRow.hidden = m !== 'replay';
    if (m === 'setup') this.tl.pause();
    this.rebuild();
  }

  private setView(v: View) {
    const d = 1.45;
    const pos: Record<View, [number, number, number]> = {
      persp: [-0.45, 0.42, 1.05],
      front: [0, 0, d],
      top: [0, d, 0.0001],
      end: [d, 0, 0],
    };
    this.camera.position.set(...pos[v]);
    this.camera.up.set(0, 1, 0);
    if (v === 'top') this.camera.up.set(0, 0, -1);
    this.controls?.target.set(0, 0, 0);
    this.camera.lookAt(0, 0, 0);
    this.controls?.update();
    this.render();
  }

  private resize() {
    if (!this.renderer) return;
    const w = this.wrap.clientWidth;
    const h = this.wrap.clientHeight;
    if (w === 0 || h === 0) return;
    this.renderer.setSize(w, h);
    this.camera.aspect = w / h;
    this.camera.updateProjectionMatrix();
    this.render();
  }

  private render() {
    this.renderer?.render(this.scene, this.camera);
  }

  // node position (beads and anchors), normalized
  private nodePos(sys: BeadSystem, node: number, out: THREE.Vector3): THREE.Vector3 {
    if (node < sys.n) out.set(sys.rest[3 * node], sys.rest[3 * node + 1], sys.rest[3 * node + 2]);
    else out.set(...sys.anchors[node - sys.n].pos);
    return out.sub(this.center).multiplyScalar(this.scale);
  }

  private rebuild() {
    const spec = this.mode === 'replay' ? this.replaySpec : this.setupSpec;
    disposeGroup(this.model);
    this.beads = this.pins = this.springs = this.hammerArrows = null;
    this.slideRings = [];
    this.pickupArrow = null;
    this.spec = spec;
    this.frames = this.mode === 'replay' && this.replaySpec ? this.replaySpec.frames : null;
    if (!spec) return this.render();
    if (this.renderer) this.renderer.setClearColor(new THREE.Color(css('--surface-1')));
    const { sys, exc, pickup } = spec;
    const n = sys.n;

    // normalize: longest extent of beads and anchors -> 1, centered
    const box = new THREE.Box3();
    const v = new THREE.Vector3();
    for (let i = 0; i < n; i++) box.expandByPoint(v.set(sys.rest[3 * i], sys.rest[3 * i + 1], sys.rest[3 * i + 2]));
    for (const a of sys.anchors) box.expandByPoint(v.set(...a.pos));
    box.getCenter(this.center);
    const size = box.getSize(new THREE.Vector3());
    this.scale = 1 / Math.max(size.x, size.y, size.z, 1e-12);

    // bead size from nearest-neighbour spacing
    let minLen = Infinity;
    const p = new THREE.Vector3();
    const q = new THREE.Vector3();
    for (let s = 0; s < sys.si.length; s++) minLen = Math.min(minLen, this.nodePos(sys, sys.si[s], p).distanceTo(this.nodePos(sys, sys.sj[s], q)));
    if (!Number.isFinite(minLen)) minLen = 0.05;
    const r = Math.min(0.32 * minLen, 0.022);

    const cBead = new THREE.Color(css('--series-1'));
    const cExc = new THREE.Color(css('--series-2'));
    const cPick = new THREE.Color(css('--series-3'));
    const cInk = new THREE.Color(css('--text-primary'));
    const cMuted = new THREE.Color(css('--text-muted'));

    // grid floor for depth cues
    const grid = new THREE.GridHelper(1.4, 14, cMuted, cMuted);
    (grid.material as THREE.Material).transparent = true;
    (grid.material as THREE.Material).opacity = 0.25;
    grid.position.y = -0.22;
    this.model.add(grid);
    this.model.add(axisTriad(new THREE.Vector3(-0.58, -0.22, -0.3), 0.08, cInk));

    // beads (pinned beads are drawn as cubes instead)
    const pickBead = (i: number) => pickup[3 * i] !== 0 || pickup[3 * i + 1] !== 0 || pickup[3 * i + 2] !== 0;
    const free: number[] = [];
    const pinned: number[] = [];
    for (let i = 0; i < n; i++) (sys.pinned[i] ? pinned : free).push(i);
    const beadMat = new THREE.MeshStandardMaterial({ roughness: 0.55, metalness: 0.1 });
    this.beads = new THREE.InstancedMesh(new THREE.SphereGeometry(r, 20, 14), beadMat, Math.max(1, n));
    this.beads.count = n;
    for (let i = 0; i < n; i++) this.beads.setColorAt(i, pickBead(i) ? cPick : cBead);
    this.model.add(this.beads);
    if (pinned.length) {
      this.pins = new THREE.InstancedMesh(new THREE.BoxGeometry(2.4 * r, 2.4 * r, 2.4 * r), new THREE.MeshStandardMaterial({ color: cInk, roughness: 0.6 }), pinned.length);
      this.model.add(this.pins);
    }
    (this.beads.userData as { free: number[]; pinned: number[] }) = { free, pinned };

    // springs
    const sg = new THREE.BufferGeometry();
    sg.setAttribute('position', new THREE.BufferAttribute(new Float32Array(sys.si.length * 6), 3));
    this.springs = new THREE.LineSegments(sg, new THREE.LineBasicMaterial({ color: cMuted }));
    this.model.add(this.springs);

    // supports
    const wallSize = 0.16;
    sys.anchors.forEach((a) => {
      const at = this.nodePos(sys, sys.n + sys.anchors.indexOf(a), new THREE.Vector3());
      const side = at.x < 0 ? -1 : 1;
      if (a.slideBead < 0) {
        const wall = new THREE.Mesh(
          new THREE.BoxGeometry(0.02, wallSize, wallSize),
          new THREE.MeshStandardMaterial({ color: cMuted, roughness: 0.8 }),
        );
        wall.position.copy(at).add(new THREE.Vector3(side * 0.01, 0, 0));
        this.model.add(wall);
      } else {
        const plane = new THREE.Mesh(
          new THREE.PlaneGeometry(wallSize, wallSize),
          new THREE.MeshBasicMaterial({ color: cMuted, transparent: true, opacity: 0.18, side: THREE.DoubleSide, depthWrite: false }),
        );
        plane.rotation.y = Math.PI / 2;
        plane.position.copy(at);
        this.model.add(plane);
        const edges = new THREE.LineSegments(new THREE.EdgesGeometry(plane.geometry), new THREE.LineBasicMaterial({ color: cMuted }));
        edges.rotation.y = Math.PI / 2;
        edges.position.copy(at);
        this.model.add(edges);
        const ring = new THREE.Mesh(new THREE.TorusGeometry(1.3 * r, 0.35 * r, 8, 24), new THREE.MeshStandardMaterial({ color: cInk }));
        ring.rotation.y = Math.PI / 2;
        ring.position.copy(at);
        this.model.add(ring);
        this.slideRings.push({ mesh: ring, bead: a.slideBead });
      }
    });

    // excitation glyphs
    const legend: string[] = [];
    const sw = (c: THREE.Color, shape = 'dot') => `<i class="glyph ${shape}" style="--c:#${c.getHexString()}"></i>`;
    legend.push(`<span>${sw(cBead)}bead</span>`);
    if (pinned.length) legend.push(`<span>${sw(cInk, 'square')}pinned bead</span>`);
    if (sys.anchors.some((a) => a.slideBead < 0)) legend.push(`<span>${sw(cMuted, 'square')}clamped end (wall)</span>`);
    if (sys.anchors.some((a) => a.slideBead >= 0)) legend.push(`<span>${sw(cInk, 'ring')}sliding end (ring on a frictionless plane)</span>`);
    const size0 = modelSize(sys);
    const dirs = (field: Float64Array) => {
      let m = 0;
      for (let i = 0; i < n; i++) m = Math.max(m, Math.hypot(field[3 * i], field[3 * i + 1], field[3 * i + 2]));
      return m;
    };
    if (this.mode === 'setup' && spec.excKind === 'pluck') {
      const m = dirs(exc.u0);
      if (m > 0) {
        const E = (GLYPH * size0) / m;
        const pts: THREE.Vector3[] = [];
        const left = sys.anchors.find((a) => a.pos[0] < sys.rest[0]);
        const right = sys.anchors.find((a) => a.pos[0] > sys.rest[3 * (n - 1)]);
        if (left) pts.push(this.nodePos(sys, n + sys.anchors.indexOf(left), new THREE.Vector3()));
        for (let i = 0; i < n; i++) {
          const w = new THREE.Vector3(exc.u0[3 * i], exc.u0[3 * i + 1], exc.u0[3 * i + 2]).multiplyScalar(E * this.scale);
          pts.push(this.nodePos(sys, i, new THREE.Vector3()).add(w));
        }
        if (right) pts.push(this.nodePos(sys, n + sys.anchors.indexOf(right), new THREE.Vector3()));
        const line = new THREE.Line(new THREE.BufferGeometry().setFromPoints(pts), new THREE.LineDashedMaterial({ color: cExc, dashSize: 0.015, gapSize: 0.01 }));
        line.computeLineDistances();
        this.model.add(line);
        let peak = 0;
        for (let i = 1; i < n; i++) if (Math.hypot(exc.u0[3 * i + 1], exc.u0[3 * i + 2]) > Math.hypot(exc.u0[3 * peak + 1], exc.u0[3 * peak + 2])) peak = i;
        const from = this.nodePos(sys, peak, new THREE.Vector3());
        const to = pts[peak + (left ? 1 : 0)];
        const len = from.distanceTo(to);
        if (len > 0) this.model.add(new THREE.ArrowHelper(to.clone().sub(from).normalize(), from, len, cExc, 0.25 * len, 0.12 * len));
        legend.push(`<span>${sw(cExc, 'dash')}pluck: initial shape, drawn ×${fmtNum(E)} (${fmtNum(m * 1000)} mm peak)</span>`);
      }
    }
    if (spec.excKind === 'strike' && this.mode === 'setup') {
      const m = dirs(exc.v0);
      if (m > 0) this.model.add(this.arrows(sys, exc.v0, m, cExc, false));
      legend.push(`<span>${sw(cExc, 'arrow')}strike: initial velocity (${fmtNum(m)} m/s peak)</span>`);
    }
    if (spec.excKind === 'hammer' && exc.force) {
      const m = dirs(exc.force.shape);
      if (m > 0) {
        this.hammerArrows = this.arrows(sys, exc.force.shape, m, cExc, true);
        this.model.add(this.hammerArrows);
      }
      legend.push(
        `<span>${sw(cExc, 'arrow')}hammer: force pulse, ${fmtNum(exc.force.F0)} N peak over ${fmtNum(exc.force.tau * 1000)} ms${this.mode === 'replay' ? ' (shown while it acts)' : ''}</span>`,
      );
    }

    // pickup
    const pk: number[] = [];
    for (let i = 0; i < n; i++) if (pickBead(i)) pk.push(i);
    if (pk.length === 1) {
      const i = pk[0];
      const d = new THREE.Vector3(pickup[3 * i], pickup[3 * i + 1], pickup[3 * i + 2]).normalize();
      const len = Math.max(0.05, 6 * r);
      const arrow = new THREE.ArrowHelper(d, this.nodePos(sys, i, new THREE.Vector3()), len, cPick, 0.3 * len, 0.18 * len);
      this.model.add(arrow);
      this.pickupArrow = { arrow, bead: i };
    }
    legend.push(`<span>${sw(cPick)}pickup: ${spec.pickupLabel}</span>`);
    this.legend.innerHTML = legend.join('');
    this.update();
  }

  // arrows proportional to a per-bead vector field; `into` = arrow ends at the bead (a push)
  private arrows(sys: BeadSystem, field: Float64Array, max: number, color: THREE.Color, into: boolean): THREE.Group {
    const g = new THREE.Group();
    for (let i = 0; i < sys.n; i++) {
      const d = new THREE.Vector3(field[3 * i], field[3 * i + 1], field[3 * i + 2]);
      const mag = d.length();
      if (mag < 0.05 * max) continue;
      const len = GLYPH * (mag / max);
      d.normalize();
      const at = this.nodePos(sys, i, new THREE.Vector3());
      const origin = into ? at.clone().addScaledVector(d, -len) : at;
      const arrow = new THREE.ArrowHelper(d, origin, len, color, Math.min(0.03, 0.3 * len), Math.min(0.016, 0.18 * len));
      arrow.userData = { bead: i, dir: d.clone(), len };
      g.add(arrow);
    }
    return g;
  }

  // position everything for the current frame (or rest, in setup)
  private update() {
    const spec = this.spec;
    if (!spec || !this.beads) return this.render();
    const { sys } = spec;
    const n = sys.n;
    const frames = this.frames;
    const off = frames ? this.tl.frame * 3 * n : 0;
    const E = this.exaggeration * this.scale;
    const P = new Array<THREE.Vector3>(n + sys.anchors.length);
    for (let i = 0; i < n; i++) {
      const p = this.nodePos(sys, i, new THREE.Vector3());
      if (frames) p.add(new THREE.Vector3(frames[off + 3 * i], frames[off + 3 * i + 1], frames[off + 3 * i + 2]).multiplyScalar(E));
      P[i] = p;
    }
    sys.anchors.forEach((a, k) => {
      const p = this.nodePos(sys, n + k, new THREE.Vector3());
      if (a.slideBead >= 0) {
        p.y = P[a.slideBead].y;
        p.z = P[a.slideBead].z;
      }
      P[n + k] = p;
    });
    const m = new THREE.Matrix4();
    const { pinned } = this.beads.userData as { free: number[]; pinned: number[] };
    for (let i = 0; i < n; i++) {
      // hide pinned beads' spheres (cubes stand in)
      m.makeScale(sys.pinned[i] ? 0 : 1, sys.pinned[i] ? 0 : 1, sys.pinned[i] ? 0 : 1).setPosition(P[i]);
      this.beads.setMatrixAt(i, m);
    }
    this.beads.instanceMatrix.needsUpdate = true;
    if (this.beads.instanceColor) this.beads.instanceColor.needsUpdate = true;
    if (this.pins) {
      pinned.forEach((i, k) => this.pins!.setMatrixAt(k, m.makeTranslation(P[i].x, P[i].y, P[i].z)));
      this.pins.instanceMatrix.needsUpdate = true;
    }
    if (this.springs) {
      const arr = (this.springs.geometry.getAttribute('position') as THREE.BufferAttribute).array as Float32Array;
      for (let s = 0; s < sys.si.length; s++) {
        const a = P[sys.si[s]];
        const b = P[sys.sj[s]];
        arr.set([a.x, a.y, a.z, b.x, b.y, b.z], 6 * s);
      }
      this.springs.geometry.getAttribute('position').needsUpdate = true;
      this.springs.geometry.computeBoundingSphere();
    }
    for (const { mesh, bead } of this.slideRings) {
      mesh.position.y = P[bead].y;
      mesh.position.z = P[bead].z;
    }
    if (this.pickupArrow) this.pickupArrow.arrow.position.copy(P[this.pickupArrow.bead]);
    if (this.hammerArrows && spec.exc.force) {
      const f = spec.exc.force;
      const t = frames ? this.tl.times[this.tl.frame] : f.tau / 2;
      const s = t >= 0 && t <= f.tau ? 0.5 * (1 - Math.cos((2 * Math.PI * t) / f.tau)) : 0;
      this.hammerArrows.visible = s > 0.02;
      for (const c of this.hammerArrows.children) {
        const { bead, dir, len } = c.userData as { bead: number; dir: THREE.Vector3; len: number };
        const l = Math.max(1e-4, len * (frames ? s : 1));
        const arrow = c as THREE.ArrowHelper;
        arrow.position.copy(P[bead]).addScaledVector(dir, -l);
        arrow.setLength(l, Math.min(0.03, 0.3 * l), Math.min(0.016, 0.18 * l));
      }
    }
    if (this.mode === 'replay' && frames) {
      this.note.textContent =
        `Displacements drawn ×${fmtNum(this.exaggeration)} (all three components, so longitudinal motion is exaggerated too). ` +
        `Drag to rotate, scroll to zoom, right-drag to pan.`;
    } else {
      this.note.textContent = 'Beads at rest. Drag to rotate, scroll to zoom, right-drag to pan.';
    }
    this.render();
  }
}

function modelSize(sys: BeadSystem): number {
  let lo = Infinity;
  let hi = -Infinity;
  for (let i = 0; i < sys.n; i++) {
    lo = Math.min(lo, sys.rest[3 * i]);
    hi = Math.max(hi, sys.rest[3 * i]);
  }
  for (const a of sys.anchors) {
    lo = Math.min(lo, a.pos[0]);
    hi = Math.max(hi, a.pos[0]);
  }
  return Math.max(hi - lo, 1e-12);
}

function textSprite(text: string, color: THREE.Color): THREE.Sprite {
  const c = document.createElement('canvas');
  c.width = c.height = 64;
  const ctx = c.getContext('2d')!;
  ctx.fillStyle = `#${color.getHexString()}`;
  ctx.font = '600 44px system-ui, sans-serif';
  ctx.textAlign = 'center';
  ctx.textBaseline = 'middle';
  ctx.fillText(text, 32, 34);
  const s = new THREE.Sprite(new THREE.SpriteMaterial({ map: new THREE.CanvasTexture(c), depthTest: false }));
  s.scale.set(0.04, 0.04, 1);
  return s;
}

function axisTriad(at: THREE.Vector3, len: number, ink: THREE.Color): THREE.Group {
  const g = new THREE.Group();
  const axes: [string, THREE.Vector3][] = [
    ['x', new THREE.Vector3(1, 0, 0)],
    ['y', new THREE.Vector3(0, 1, 0)],
    ['z', new THREE.Vector3(0, 0, 1)],
  ];
  for (const [name, d] of axes) {
    g.add(new THREE.ArrowHelper(d, at, len, ink, 0.2 * len, 0.1 * len));
    const label = textSprite(name, ink);
    label.position.copy(at).addScaledVector(d, len * 1.3);
    g.add(label);
  }
  return g;
}

function disposeGroup(g: THREE.Group) {
  g.traverse((o) => {
    const m = o as THREE.Mesh;
    m.geometry?.dispose();
    const mat = m.material as THREE.Material | THREE.Material[] | undefined;
    if (Array.isArray(mat)) mat.forEach((x) => x.dispose());
    else mat?.dispose();
  });
  g.clear();
}