SteelXR2/src/components/three/XRControllerMeasure.tsx

import { useRef, useState } from 'react';
import { useFrame, useThree } from '@react-three/fiber';
import * as THREE from 'three';
import { useModelStore } from '@/stores/useModelStore';
import { resolveSnap, detectCircularEdgeAtPoint, findNearestEdgeSegment } from './SmartMeasure';

const TRIG_ON = 0.7;
const TRIG_OFF = 0.3;
const BTN_ON = 0.6;
const MAX_RAY = 10; // meters

/**
 * Right-controller trigger driven measurement for AR.
 *
 * Mapping (right hand):
 *   - Trigger  → add measurement point (with snap)
 *   - Button A (gp.buttons[4]) → undo last point/measurement
 *   - Button B (gp.buttons[5]) → clear all measurements
 * Left hand:
 *   - Trigger → toggle vertex snap ON/OFF
 *
 * Renders a thin laser from the right controller to the snap candidate
 * (green if snapped to vertex/edge, amber if surface).
 */
export function XRControllerMeasure() {
  const { scene, gl } = useThree();
  const trigState = useRef(false);
  const aState = useRef(false);
  const bState = useRef(false);
  const lTrigState = useRef(false);
  const [snapEnabled, setSnapEnabled] = useState(true);

  const raycaster = useRef(new THREE.Raycaster());
  const tmpOrigin = useRef(new THREE.Vector3());
  const tmpDir = useRef(new THREE.Vector3());
  const tmpQuat = useRef(new THREE.Quaternion());

  // Visual laser refs
  const laserRef = useRef<THREE.Line>(null);
  const tipRef = useRef<THREE.Mesh>(null);
  const laserGeom = useRef<THREE.BufferGeometry>(
    new THREE.BufferGeometry().setFromPoints([new THREE.Vector3(), new THREE.Vector3(0, 0, -1)])
  );
  const laserMat = useRef(new THREE.LineBasicMaterial({ color: '#22c55e', transparent: true, opacity: 0.7, depthTest: false }));
  const tipColor = useRef(new THREE.Color('#22c55e'));

  useFrame((_state, _dt, frame: XRFrame | undefined) => {
    const measureMode = useModelStore.getState().measureMode;
    if (laserRef.current) laserRef.current.visible = false;
    if (tipRef.current) tipRef.current.visible = false;
    if (!measureMode || !frame) return;

    const session = frame.session;
    const refSpace = gl.xr.getReferenceSpace();
    if (!session || !refSpace) return;

    let right: XRInputSource | null = null;
    let left: XRInputSource | null = null;
    for (const src of session.inputSources) {
      if (src.handedness === 'right') right = src;
      else if (src.handedness === 'left') left = src;
    }

    // ── Left trigger: toggle snap ─────────────────────────────────────
    if (left) {
      const v = left.gamepad?.buttons?.[0]?.value ?? (left.gamepad?.buttons?.[0]?.pressed ? 1 : 0);
      if (!lTrigState.current && v > TRIG_ON) {
        lTrigState.current = true;
        setSnapEnabled((s) => !s);
      } else if (lTrigState.current && v < TRIG_OFF) {
        lTrigState.current = false;
      }
    }

    if (!right) return;
    const gp = right.gamepad;

    // ── Compute right-controller ray ──────────────────────────────────
    const raySpace = right.targetRaySpace ?? right.gripSpace;
    if (!raySpace) return;
    const pose = frame.getPose(raySpace, refSpace);
    if (!pose) return;

    const m = new THREE.Matrix4().fromArray(pose.transform.matrix);
    tmpOrigin.current.setFromMatrixPosition(m);
    tmpQuat.current.setFromRotationMatrix(m);
    tmpDir.current.set(0, 0, -1).applyQuaternion(tmpQuat.current).normalize();

    raycaster.current.set(tmpOrigin.current, tmpDir.current);
    raycaster.current.far = MAX_RAY;

    const hits = raycaster.current.intersectObjects(scene.children, true);
    const hit = hits.find((h) => {
      const o = h.object;
      if (!(o instanceof THREE.Mesh)) return false;
      if (o.userData.__edgeLine) return false;
      if (o.geometry instanceof THREE.SphereGeometry) return false;
      if (o.geometry instanceof THREE.RingGeometry) return false;
      if (o.geometry instanceof THREE.PlaneGeometry) return false;
      return true;
    });

    let snappedPoint: THREE.Vector3 | null = null;
    let snapKind: 'vertex' | 'edge' | 'surface' = 'surface';

    if (hit && hit.object instanceof THREE.Mesh) {
      if (snapEnabled) {
        // Use a virtual canvas size based on session — fall back to current renderer size
        const size = gl.getSize(new THREE.Vector2());
        const fakeCam = new THREE.PerspectiveCamera(60, 1, 0.01, 100);
        fakeCam.position.copy(tmpOrigin.current);
        fakeCam.quaternion.copy(tmpQuat.current);
        fakeCam.updateMatrixWorld(true);
        const snap = resolveSnap(
          hit.object,
          hit.point,
          fakeCam,
          { width: size.x || 1024, height: size.y || 1024 },
          14,
          18
        );
        snappedPoint = snap.point;
        snapKind = snap.type;
      } else {
        snappedPoint = hit.point.clone();
        snapKind = 'surface';
      }
    }

    // ── Update laser visual ───────────────────────────────────────────
    if (laserRef.current && tipRef.current) {
      const end = snappedPoint ?? tmpOrigin.current.clone().add(tmpDir.current.clone().multiplyScalar(MAX_RAY));
      const positions = laserGeom.current.attributes.position as THREE.BufferAttribute;
      positions.setXYZ(0, tmpOrigin.current.x, tmpOrigin.current.y, tmpOrigin.current.z);
      positions.setXYZ(1, end.x, end.y, end.z);
      positions.needsUpdate = true;
      laserGeom.current.computeBoundingSphere();

      const color = snapKind === 'vertex' ? '#22c55e' : snapKind === 'edge' ? '#3b82f6' : '#eab308';
      tipColor.current.set(color);
      (laserRef.current.material as THREE.LineBasicMaterial).color.set(color);
      ((tipRef.current.material as THREE.MeshBasicMaterial)).color.copy(tipColor.current);

      laserRef.current.visible = true;
      if (snappedPoint) {
        tipRef.current.visible = true;
        tipRef.current.position.copy(snappedPoint);
        // Scale tip with distance from controller (~12mm at 1m)
        const dist = tmpOrigin.current.distanceTo(snappedPoint);
        const s = Math.max(0.004, dist * 0.012);
        tipRef.current.scale.setScalar(s);
      }
    }

    // ── Right trigger: add point ──────────────────────────────────────
    const trigVal = gp?.buttons?.[0]?.value ?? (gp?.buttons?.[0]?.pressed ? 1 : 0);
    if (!trigState.current && trigVal > TRIG_ON) {
      trigState.current = true;
      if (snappedPoint) {
        useModelStore.getState().addMeasurePoint({
          x: snappedPoint.x, y: snappedPoint.y, z: snappedPoint.z,
        });
      }
    } else if (trigState.current && trigVal < TRIG_OFF) {
      trigState.current = false;
    }

    // ── Button A (undo) ──────────────────────────────────────────────
    const aBtn = gp?.buttons?.[4];
    const aVal = aBtn ? (aBtn.value || (aBtn.pressed ? 1 : 0)) : 0;
    if (!aState.current && aVal > BTN_ON) {
      aState.current = true;
      useModelStore.getState().undoLastMeasurement();
    } else if (aState.current && aVal < TRIG_OFF) {
      aState.current = false;
    }

    // ── Button B (clear) ─────────────────────────────────────────────
    const bBtn = gp?.buttons?.[5];
    const bVal = bBtn ? (bBtn.value || (bBtn.pressed ? 1 : 0)) : 0;
    if (!bState.current && bVal > BTN_ON) {
      bState.current = true;
      useModelStore.getState().clearMeasurements();
    } else if (bState.current && bVal < TRIG_OFF) {
      bState.current = false;
    }
  });

  // Construct the laser Line object once so we can attach via <primitive>
  const laserObject = useRef<THREE.Line>();
  if (!laserObject.current) {
    laserObject.current = new THREE.Line(laserGeom.current, laserMat.current);
    laserObject.current.frustumCulled = false;
    laserObject.current.renderOrder = 999;
    (laserRef as React.MutableRefObject<THREE.Line | null>).current = laserObject.current;
  }

  return (
    <>
      <primitive object={laserObject.current} />
      <mesh ref={tipRef} frustumCulled={false} renderOrder={999}>
        <sphereGeometry args={[1, 12, 12]} />
        <meshBasicMaterial color="#22c55e" depthTest={false} transparent opacity={0.9} />
      </mesh>
    </>
  );
}