/* Newbury Exhibit — the Hidden Side (multi-viewer AR)
 *
 * This is the payoff: point a phone at any Newbury Crest and the ghosts appear
 * anchored to the physical town, moving along their authored paths, hidden behind
 * buildings — and every phone in the room sees the SAME ghosts in the SAME spots,
 * each from its own viewpoint.
 *
 * How the shared world works:
 *   1. Camera detects a marker -> POS-IT gives camera pose RELATIVE TO THAT MARKER.
 *   2. scene.json says where that marker sits in WORLD coords (anchor position/yaw).
 *   3. Compose: world->camera transform. Any device seeing any known marker lands
 *      in the same world frame, so world-space ghosts line up for everyone.
 *   4. The server streams ghost world-positions at 20Hz; we place them in the world
 *      group, which is transformed into camera space. Buildings are invisible
 *      depth-only occluders so ghosts vanish behind them.
 */
import * as THREE from 'three';
import { buildLegoHeadGhost } from './lego-head-ghost.js?v=1';
import { TunedDetector, MarkerTracker } from './detect-tuned.js?v=1';

const MM = 0.001; // world mm -> three.js metres

const $ = (id) => document.getElementById(id);
const video = $('video');
const canvas = $('three');
const startScreen = $('start'), goBtn = $('go'), errEl = $('err');
const hud = $('hud'), countEl = $('count'), lockTag = $('lock'), line2 = $('line2');
const ghostNEl = $('ghostN'), realignEl = $('realign');

const grab = document.createElement('canvas');
const gctx = grab.getContext('2d', { willReadFrequently: true });

let detector = null, tracker = null, focalPx = 0, running = false;

// ---- Three.js scene -------------------------------------------------------
const renderer = new THREE.WebGLRenderer({ canvas, alpha: true, antialias: true });
renderer.setClearColor(0x000000, 0);
renderer.sortObjects = true;

const scene3d = new THREE.Scene();
scene3d.add(new THREE.AmbientLight(0xb8c4e0, 1.15));
const key = new THREE.DirectionalLight(0xffffff, 0.75);
key.position.set(0.3, 1, 0.5); scene3d.add(key);

// Camera sits at origin looking down -Z. We move the WORLD, not the camera.
let camera = new THREE.PerspectiveCamera(55, 1, 5, 8000); // mm units (large far)

// `world` holds everything in build coordinates; its matrix = world->camera.
const world = new THREE.Group();
world.matrixAutoUpdate = false;
scene3d.add(world);

// sub-groups inside world
const gGhosts = new THREE.Group(); world.add(gGhosts);
const gOccluders = new THREE.Group(); world.add(gOccluders);
const gAnchorViz = new THREE.Group(); world.add(gAnchorViz); // subtle anchor markers

// ---- Scene data from server ----------------------------------------------
let sceneMeta = null;         // { world, anchors[], blockers[] }
const anchorById = new Map(); // fiducialId -> anchor (world placement)
const ghostObjs = new Map();  // spawnId -> { group, targetPos, lastColor }

// ---- Marker + pose --------------------------------------------------------
const MARKER_SIZE_MM = 96;    // outer black square; must match printed/built markers
let posit = null;

function centredCorners(m, w, h) {
  return m.corners.map((c) => ({ x: c.x - w / 2, y: h / 2 - c.y }));
}

// Build the marker->camera matrix (POS-IT), converting js-aruco axes to three.js.
// js-aruco POS-IT gives the marker pose in a camera frame with X right, Y up,
// Z toward viewer. Three.js's camera looks down -Z with Y up, so we convert by
// premultiplying F = diag(1,-1,-1) (flip Y and Z). This is the SAME conversion
// the working ar-head.js uses, and it is correct for pitch AND roll: tilting the
// phone up moves an anchored ghost down-screen, and rolling the phone rotates the
// ghost the same way. (An earlier F*M*F "fix" inverted both — don't reintroduce it.)
const flipYZ = new THREE.Matrix4().makeScale(1, -1, -1);

// Convert one POS-IT (R,t) solution into a three.js marker->camera matrix.
function poseToMatrix(R, t) {
  const M = new THREE.Matrix4().set(
    R[0][0], R[0][1], R[0][2], t[0],
    R[1][0], R[1][1], R[1][2], t[1],
    R[2][0], R[2][1], R[2][2], t[2],
    0, 0, 0, 1
  );
  M.premultiply(flipYZ);  // F * M  (flip Y & Z: POS-IT frame -> three.js frame)
  return M;
}

// POS-IT returns TWO solutions for a planar marker (the real pose and a mirror
// twin). Up close they have near-equal error and the solver flip-flops between
// them -> "rotate 90deg, ghost spins 180deg". We return both; alignWorld picks
// the one whose resulting world orientation is closest to the previous frame.
function markerToCameraMatrix(m) {
  const pose = posit.pose(centredCorners(m, grab.width, grab.height));
  if (!pose) return null;
  const best = { M: poseToMatrix(pose.bestRotation, pose.bestTranslation), err: pose.bestError };
  let alt = null;
  if (pose.alternativeRotation && pose.alternativeTranslation) {
    alt = { M: poseToMatrix(pose.alternativeRotation, pose.alternativeTranslation), err: pose.alternativeError };
  }
  const t = pose.bestTranslation;
  return { best, alt, dist: Math.hypot(t[0], t[1], t[2]) };
}

// anchor world placement -> matrix mapping the MARKER's local frame (as POS-IT
// sees it: marker in its XY plane, +Z out of the printed face) into WORLD space.
//
// A crest lying flat on the table has its face pointing UP, so the marker's
// local +Z must map to world +Y. That's a -90 deg rotation about X. On top of
// that we apply the anchor's yaw (spin on the table) and its world position.
//
//   mount 'flat' (default): marker lies on the table, face up   -> rotX(-90) then yaw about Y
//   mount 'wall':           marker stands vertical, face outward -> yaw about Y only
//
// The order matters: worldFromMarker = T(pos) * Ry(yaw) * planeTilt.
function anchorToWorldMatrix(anchor) {
  const p = anchor.position || { x: 0, y: 0, z: 0 };
  const yaw = THREE.MathUtils.degToRad(anchor.rotationDeg?.y || 0);
  const mount = anchor.mount || 'flat';

  const planeTilt = new THREE.Matrix4();
  if (mount === 'wall') planeTilt.identity();          // face already points sideways
  else planeTilt.makeRotationX(-Math.PI / 2);           // flat: tip face up to +Y

  const spin = new THREE.Matrix4().makeRotationY(yaw);
  const trans = new THREE.Matrix4().setPosition(p.x, p.y, p.z);

  // trans * spin * planeTilt
  return trans.multiply(spin).multiply(planeTilt);
}

// Given a detected marker with a known anchor, compute world->camera and apply
// it to the `world` group so everything in world coords renders correctly.
//
// Planar-ambiguity defeat: POS-IT gives two candidate poses. We compute the
// resulting world->camera for each, and keep whichever is closest to LAST
// frame's transform (temporal consistency). This stops the mirror-twin flip
// that made the ghost spin the wrong way. On first lock we take `best`.
const tmpAnchorInv = new THREE.Matrix4();
let lastWorldToCamera = null; // previous frame's chosen transform (for continuity)

function worldToCameraFor(mMatrix, anchorInv) {
  return new THREE.Matrix4().multiplyMatrices(mMatrix, anchorInv);
}

// crude distance between two 4x4s: sum of squared element differences of the
// rotation part (enough to tell the real pose from its mirror twin).
function matDist(a, b) {
  const ea = a.elements, eb = b.elements;
  let s = 0;
  for (const i of [0, 1, 2, 4, 5, 6, 8, 9, 10]) { const d = ea[i] - eb[i]; s += d * d; }
  return s;
}

function alignWorld(m2c, anchor) {
  const aToW = anchorToWorldMatrix(anchor);
  tmpAnchorInv.copy(aToW).invert();

  const candBest = worldToCameraFor(m2c.best.M, tmpAnchorInv);
  let chosen = candBest;

  if (m2c.alt) {
    const candAlt = worldToCameraFor(m2c.alt.M, tmpAnchorInv);
    if (lastWorldToCamera) {
      // pick the candidate closest to where the world was last frame
      chosen = matDist(candAlt, lastWorldToCamera) < matDist(candBest, lastWorldToCamera)
        ? candAlt : candBest;
    } else {
      // first lock: trust the lower-error (best) solution
      chosen = candBest;
    }
  }

  world.matrix.copy(chosen);
  world.matrixWorldNeedsUpdate = true;
  lastWorldToCamera = chosen.clone();
}

// ---- Multi-marker fusion --------------------------------------------------
// When several crests are visible at once, each yields its own world->camera
// estimate. Picking a single "winner" makes the world snap as detection flickers
// between them. Because every crest's world position is known and the rig is
// rigid, we instead FUSE all visible estimates into one consensus pose:
//   * for each marker, resolve its two POS-IT solutions to the one consistent
//     with last frame (defeats the planar flip), producing one world->camera;
//   * weight each by confidence (bigger, lower-error, non-coasting = more trust);
//   * average translations linearly and rotations via quaternion slerp.
// Result: smoother and more accurate than any single marker, with no snap.
const _p = new THREE.Vector3(), _q = new THREE.Quaternion(), _s = new THREE.Vector3();

function resolveMarkerWorldToCamera(m2c, anchor) {
  // returns the single world->camera for this marker, flip-resolved vs last frame
  const aToW = anchorToWorldMatrix(anchor);
  const inv = new THREE.Matrix4().copy(aToW).invert();
  const candBest = worldToCameraFor(m2c.best.M, inv);
  if (!m2c.alt) return candBest;
  const candAlt = worldToCameraFor(m2c.alt.M, inv);
  if (!lastWorldToCamera) {
    return (m2c.alt.err < m2c.best.err) ? candBest : candBest; // first lock: best
  }
  return matDist(candAlt, lastWorldToCamera) < matDist(candBest, lastWorldToCamera)
    ? candAlt : candBest;
}

function fuseWorld(observations) {
  // observations: [{ m2c, anchor, weight }]
  if (observations.length === 0) return false;

  // Resolve each to a single world->camera, decompose to pos+quat.
  const parts = [];
  for (const o of observations) {
    const wc = resolveMarkerWorldToCamera(o.m2c, o.anchor);
    const p = new THREE.Vector3(), q = new THREE.Quaternion(), s = new THREE.Vector3();
    wc.decompose(p, q, s);
    parts.push({ p, q, w: o.weight });
  }

  // Weighted translation average.
  const totalW = parts.reduce((a, b) => a + b.w, 0) || 1;
  const pos = new THREE.Vector3();
  for (const pt of parts) pos.addScaledVector(pt.p, pt.w / totalW);

  // Weighted rotation average via incremental slerp. Keep all quats on the same
  // hemisphere as the first (q and -q are the same rotation; slerp needs care).
  const q0 = parts[0].q.clone();
  let quat = q0.clone();
  let accum = parts[0].w;
  for (let i = 1; i < parts.length; i++) {
    let qi = parts[i].q.clone();
    if (q0.dot(qi) < 0) qi.set(-qi.x, -qi.y, -qi.z, -qi.w); // hemisphere align
    accum += parts[i].w;
    quat.slerp(qi, parts[i].w / accum); // incremental weighted mean
    quat.normalize();
  }

  const fused = new THREE.Matrix4().compose(pos, quat, new THREE.Vector3(1, 1, 1));
  world.matrix.copy(fused);
  world.matrixWorldNeedsUpdate = true;
  lastWorldToCamera = fused.clone();
  return true;
}

// ---- Build occluders + anchor viz once scene arrives ----------------------
function buildStaticWorld() {
  // occluders: invisible depth-only boxes so ghosts hide behind buildings
  gOccluders.clear();
  for (const bl of (sceneMeta.blockers || [])) {
    const sz = bl.size;
    const geo = new THREE.BoxGeometry(sz.x, sz.y, sz.z);
    const mat = new THREE.MeshBasicMaterial({ colorWrite: false }); // writes depth only
    const mesh = new THREE.Mesh(geo, mat);
    mesh.position.set(bl.position.x + sz.x / 2, bl.position.y + sz.y / 2, bl.position.z + sz.z / 2);
    mesh.renderOrder = -1; // draw occluders before ghosts so depth is ready
    gOccluders.add(mesh);
  }

  // very subtle anchor dots so you can see the crests are recognised (optional)
  gAnchorViz.clear();
  for (const a of (sceneMeta.anchors || [])) {
    const dot = new THREE.Mesh(
      new THREE.SphereGeometry(6, 10, 10),
      new THREE.MeshBasicMaterial({ color: 0x00e5c0, transparent: true, opacity: 0.25 })
    );
    dot.position.set(a.position.x, a.position.y, a.position.z);
    gAnchorViz.add(dot);
  }
}

// ---- Ghosts ---------------------------------------------------------------
const LURE_KEY = { Red: 'Red', Yellow: 'Yellow', Blue: 'Blue' };

function ensureGhost(g) {
  let obj = ghostObjs.get(g.spawnId);
  if (!obj) {
    const colorKey = colorForGhost(g);
    const group = buildLegoHeadGhost(THREE, colorKey);
    group.scale.setScalar(1.0); // ghost built in mm already
    gGhosts.add(group);
    obj = { group, targetPos: new THREE.Vector3(), lastColor: colorKey };
    ghostObjs.set(g.spawnId, obj);
  }
  return obj;
}

// The server sends ghostId; we map to a lure colour via the roster if we have it.
let roster = null;
function colorForGhost(g) {
  if (g.color) return LURE_KEY[g.color] || 'Blue';
  if (roster && g.ghostId && roster.has(g.ghostId)) return roster.get(g.ghostId).color || 'Blue';
  return 'Blue';
}

function updateGhosts(list, dtServer) {
  const seen = new Set();
  for (const g of list) {
    seen.add(g.spawnId);
    const obj = ensureGhost(g);
    // server pos is world mm; set as smoothing target
    obj.targetPos.set(g.pos.x, g.pos.y, g.pos.z);
    obj._yaw = THREE.MathUtils.degToRad(g.yawDeg || 0);
  }
  // remove ghosts no longer present
  for (const [id, obj] of ghostObjs) {
    if (!seen.has(id)) { gGhosts.remove(obj.group); ghostObjs.delete(id); }
  }
  ghostNEl.textContent = ghostObjs.size;
}

// ---- WebSocket sync -------------------------------------------------------
let ws = null, wsConnected = false, lastState = 0;

function connectWS() {
  const proto = location.protocol === 'https:' ? 'wss' : 'ws';
  ws = new WebSocket(`${proto}://${location.host}/sync`);
  ws.onopen = () => { wsConnected = true; setLine(); };
  ws.onclose = () => { wsConnected = false; setLine(); setTimeout(connectWS, 1500); };
  ws.onerror = () => { /* onclose will handle reconnect */ };
  ws.onmessage = (ev) => {
    let msg; try { msg = JSON.parse(ev.data); } catch (_) { return; }
    if (msg.type === 'hello') {
      sceneMeta = { world: msg.world, anchors: msg.anchors || [], blockers: msg.blockers || [] };
      anchorById.clear();
      for (const a of sceneMeta.anchors) anchorById.set(a.fiducialId, a);
      buildStaticWorld();
      setLine();
    } else if (msg.type === 'state') {
      lastState = performance.now();
      updateGhosts(msg.ghosts, msg.serverTime);
    }
  };
}

// fetch roster once for colour mapping (optional, best-effort)
async function loadRoster() {
  try {
    const r = await fetch('/api/ghosts');
    const list = await r.json();
    roster = new Map(list.map((g) => [g.id, g]));
  } catch (_) { roster = null; }
}

// ---- Alignment state ------------------------------------------------------
let aligned = false;
let lastAlign = 0;
const REALIGN_GRACE_MS = 1500; // keep last alignment briefly after marker lost

function setLock(state) {
  // state: 'locked' | 'holding' | 'searching'
  if (state === 'locked') { lockTag.textContent = 'locked'; lockTag.className = 'tag on'; realignEl.classList.remove('show'); }
  else if (state === 'holding') { lockTag.textContent = 'holding'; lockTag.className = 'tag warn'; }
  else { lockTag.textContent = 'aligning…'; lockTag.className = 'tag off'; }
}
function setLine() {
  if (!wsConnected) { line2.textContent = 'connecting to exhibit…'; return; }
  if (!sceneMeta) { line2.textContent = 'loading scene…'; return; }
  if (!aligned) { line2.textContent = 'point at a Newbury Crest'; return; }
  line2.innerHTML = `aligned · <b>${ghostObjs.size}</b> ghosts live`;
}

// ---- Camera + sizing ------------------------------------------------------
function sizeAll() {
  const vw = video.videoWidth, vh = video.videoHeight;
  if (!vw || !vh) return;
  grab.width = vw; grab.height = vh;
  const scale = Math.max(window.innerWidth / vw, window.innerHeight / vh);
  const dw = vw * scale, dh = vh * scale;
  for (const el of [video, canvas]) { el.style.width = dw + 'px'; el.style.height = dh + 'px'; }
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
  renderer.setSize(dw, dh, false);
  focalPx = vw;
  const fovY = 2 * Math.atan((vh / 2) / focalPx) * (180 / Math.PI);
  camera.fov = fovY; camera.aspect = dw / dh; camera.updateProjectionMatrix();
  posit = new POS.Posit(MARKER_SIZE_MM, focalPx);
}

async function start() {
  errEl.textContent = '';
  try {
    const stream = await navigator.mediaDevices.getUserMedia({
      video: { facingMode: { ideal: 'environment' }, width: { ideal: 1280 }, height: { ideal: 720 } }, audio: false,
    });
    video.srcObject = stream; await video.play();
  } catch (e) {
    errEl.textContent = 'Camera unavailable: ' + (e.message || e.name) + '. iOS needs HTTPS + camera permission.';
    return;
  }
  detector = new TunedDetector({ dictionaryName: 'ARUCO_4X4_1000' }).setPreset('forgiving');
  tracker = new MarkerTracker({ coastMs: 120, smooth: true }); // short coast: don't linger on stale pose while panning

  await new Promise((res) => { if (video.readyState >= 2) res(); else video.onloadeddata = () => res(); });
  sizeAll();
  window.addEventListener('resize', sizeAll);

  await loadRoster();
  connectWS();

  startScreen.classList.add('hidden');
  hud.classList.remove('hidden'); countEl.classList.remove('hidden');
  running = true;
  requestAnimationFrame(loop);
}

// ---- Main loop ------------------------------------------------------------
function loop(now) {
  if (!running) return;
  requestAnimationFrame(loop);
  if (video.readyState < 2 || !grab.width) return;

  // 1) detect + track markers
  gctx.drawImage(video, 0, 0, grab.width, grab.height);
  const img = gctx.getImageData(0, 0, grab.width, grab.height);
  let raw = [];
  try { raw = detector.detect(img, { width: grab.width, height: grab.height }); } catch (_) { raw = []; }
  const markers = tracker.update(raw, now);

  // 2) Fuse ALL visible KNOWN markers into one consensus world pose.
  //    Each contributes an estimate weighted by confidence (area, low pose
  //    error, not coasting). Fusing removes the snap you get when a single
  //    "winner" marker flickers between two visible crests, and is more
  //    accurate because the rig geometry is known and rigid.
  const observations = [];
  let anyLive = false;
  for (const m of markers) {
    if (!anchorById.has(m.id)) continue;
    const m2c = markerToCameraMatrix(m);
    if (!m2c) continue;
    // confidence: lower pose error + live (not coasting) + nearer = more trust
    const errW = 1 / (1 + (m2c.best.err ?? 0));       // lower error -> higher weight
    const coastW = m.coasting ? 0.35 : 1;             // trust live reads more
    const nearW = 1 / (1 + m2c.dist / 1000);          // nearer markers slightly favoured
    const weight = errW * coastW * nearW;
    observations.push({ m2c, anchor: anchorById.get(m.id), weight });
    if (!m.coasting) anyLive = true;
  }

  if (observations.length > 0) {
    fuseWorld(observations);
    aligned = true; lastAlign = now;
    world.visible = true;
    setLock(anyLive ? 'locked' : 'holding');
  } else if (aligned && now - lastAlign < REALIGN_GRACE_MS) {
    // keep showing world briefly using last alignment
    setLock('holding');
  } else {
    aligned = false;
    world.visible = false;
    setLock('searching');
    if (sceneMeta && wsConnected) realignEl.classList.add('show');
  }
  setLine();

  // 3) animate ghosts: smooth toward server target, apply yaw + idle wobble
  const t = now / 1000;
  for (const obj of ghostObjs.values()) {
    obj.group.position.lerp(obj.targetPos, 0.5); // smooth network updates (snappier follow)
    if (obj._yaw != null) {
      // face travel/authored yaw, eased
      const cur = obj.group.rotation.y;
      let d = obj._yaw - cur; while (d > Math.PI) d -= 2 * Math.PI; while (d < -Math.PI) d += 2 * Math.PI;
      obj.group.rotation.y = cur + d * 0.15;
    }
    obj.group.userData.update?.(t);
  }

  // 4) stale-state watchdog: if server went quiet, dim the count
  if (wsConnected && now - lastState > 3000) ghostNEl.style.opacity = '0.4';
  else ghostNEl.style.opacity = '1';

  renderer.render(scene3d, camera);
}

goBtn.addEventListener('click', start);