mattercontrol/src/composables/useTranslateGizmo.ts

import * as THREE from 'three'

export interface TranslateGizmoOptions {
  arrowLength?: number
  arrowHeadLength?: number
  arrowHeadRadius?: number
  shaftRadius?: number
  screenSize?: number // Target screen size in pixels
}

export interface TranslateGizmo {
  group: THREE.Group
  show: () => void
  hide: () => void
  setPositionAndBounds: (center: THREE.Vector3, box: THREE.Box3) => void
  getHoveredAxis: (raycaster: THREE.Raycaster) => 'x' | 'y' | 'z' | null
  setHighlightedAxis: (axis: 'x' | 'y' | 'z' | null) => void
  updateScale: (camera: THREE.Camera) => void
  dispose: () => void
}

const AXIS_COLORS = {
  x: 0xe74c3c, // Red
  y: 0x2ecc71, // Green
  z: 0x3498db, // Blue
}

const HIGHLIGHT_COLORS = {
  x: 0xff6b6b,
  y: 0x5dff7f,
  z: 0x5dafff,
}

export function createTranslateGizmo(options: TranslateGizmoOptions = {}): TranslateGizmo {
  const {
    arrowLength = 1.0, // Extension past bounding box
    arrowHeadLength = 0.2,
    arrowHeadRadius = 0.08,
    shaftRadius = 0.025,
    screenSize = 150, // Target size in screen pixels
  } = options

  const group = new THREE.Group()
  group.name = '__translateGizmo__'
  group.visible = false

  // Store materials for highlight toggling
  const materials: Record<string, THREE.MeshBasicMaterial> = {}

  // Store arrow references for positioning
  const arrows: Record<string, THREE.Group> = {}

  // Store shaft and cone meshes for dynamic length adjustment
  const shafts: Record<string, THREE.Mesh> = {}
  const cones: Record<string, THREE.Mesh> = {}

  // Store world-space half-sizes from center to each bounding box face
  type ArrowKey = 'x_pos' | 'x_neg' | 'y_pos' | 'y_neg' | 'z_pos' | 'z_neg'
  const boundsHalfSizes: Record<ArrowKey, number> = {
    x_pos: 0,
    x_neg: 0,
    y_pos: 0,
    y_neg: 0,
    z_pos: 0,
    z_neg: 0,
  }

  let currentScale = 1

  function createArrow(axis: 'x' | 'y' | 'z', negative: boolean): THREE.Group {
    const arrowGroup = new THREE.Group()
    const suffix = negative ? '_neg' : '_pos'
    arrowGroup.name = `__gizmo_${axis}${suffix}__`
    arrowGroup.userData = { axis, isGizmoHandle: true, negative }

    const color = AXIS_COLORS[axis]

    // Create shaft with unit height - will be scaled dynamically
    const shaftGeometry = new THREE.CylinderGeometry(shaftRadius, shaftRadius, 1, 8)
    const shaftMaterial = new THREE.MeshBasicMaterial({ color })
    materials[`${axis}${suffix}_shaft`] = shaftMaterial
    const shaft = new THREE.Mesh(shaftGeometry, shaftMaterial)
    shaft.userData = { axis, isGizmoHandle: true, negative }
    shafts[`${axis}${suffix}`] = shaft

    // Create cone (arrow head) - positioned dynamically
    const coneGeometry = new THREE.ConeGeometry(arrowHeadRadius, arrowHeadLength, 16)
    const coneMaterial = new THREE.MeshBasicMaterial({ color })
    materials[`${axis}${suffix}_cone`] = coneMaterial
    const cone = new THREE.Mesh(coneGeometry, coneMaterial)
    cone.userData = { axis, isGizmoHandle: true, negative }
    cones[`${axis}${suffix}`] = cone

    arrowGroup.add(shaft)
    arrowGroup.add(cone)

    // Rotate to point in correct direction
    if (axis === 'x') {
      arrowGroup.rotation.z = negative ? Math.PI / 2 : -Math.PI / 2
    } else if (axis === 'y') {
      arrowGroup.rotation.z = negative ? Math.PI : 0
    } else {
      // z axis
      arrowGroup.rotation.x = negative ? -Math.PI / 2 : Math.PI / 2
    }

    arrows[`${axis}${suffix}`] = arrowGroup
    return arrowGroup
  }

  // Create arrows in both directions for each axis
  group.add(createArrow('x', false))
  group.add(createArrow('x', true))
  group.add(createArrow('y', false))
  group.add(createArrow('y', true))
  group.add(createArrow('z', false))
  group.add(createArrow('z', true))

  function show(): void {
    group.visible = true
  }

  function hide(): void {
    group.visible = false
  }

  function setPositionAndBounds(center: THREE.Vector3, box: THREE.Box3): void {
    group.position.copy(center)

    // Store world-space distances from center to each bounding box face
    boundsHalfSizes.x_pos = box.max.x - center.x
    boundsHalfSizes.x_neg = center.x - box.min.x
    boundsHalfSizes.y_pos = box.max.y - center.y
    boundsHalfSizes.y_neg = center.y - box.min.y
    boundsHalfSizes.z_pos = box.max.z - center.z
    boundsHalfSizes.z_neg = center.z - box.min.z

    // Update arrow lengths based on bounds
    updateArrowLengths()
  }

  function updateArrowLengths(): void {
    // Arrows originate from center and extend past bounding box by arrowLength
    // Total length = boundsHalfSize + arrowLength (in world space)
    // In local space, divide by currentScale
    const keys: ArrowKey[] = ['x_pos', 'x_neg', 'y_pos', 'y_neg', 'z_pos', 'z_neg']

    for (const key of keys) {
      const shaft = shafts[key]
      const cone = cones[key]
      if (!shaft || !cone) continue

      // Total world-space length from center to arrow tip
      const worldLength = boundsHalfSizes[key] + arrowLength
      // Convert to local space (compensating for group scale)
      const localLength = worldLength / currentScale

      // Shaft length = total length - cone head length
      const shaftLength = localLength - arrowHeadLength

      // Update shaft: scale Y to desired length, position at midpoint
      shaft.scale.y = Math.max(0.01, shaftLength)
      shaft.position.y = shaftLength / 2

      // Position cone at end of shaft
      cone.position.y = shaftLength + arrowHeadLength / 2
    }
  }

  function getHoveredAxis(raycaster: THREE.Raycaster): 'x' | 'y' | 'z' | null {
    if (!group.visible) return null

    const intersects = raycaster.intersectObject(group, true)
    if (intersects.length > 0) {
      const hit = intersects[0]!.object
      if (hit.userData?.axis) {
        return hit.userData.axis as 'x' | 'y' | 'z'
      }
    }
    return null
  }

  function setHighlightedAxis(axis: 'x' | 'y' | 'z' | null): void {
    // Reset all to original colors
    for (const ax of ['x', 'y', 'z'] as const) {
      const color = ax === axis ? HIGHLIGHT_COLORS[ax] : AXIS_COLORS[ax]
      for (const suffix of ['_pos', '_neg']) {
        materials[`${ax}${suffix}_shaft`]?.color.setHex(color)
        materials[`${ax}${suffix}_cone`]?.color.setHex(color)
      }
    }
  }

  function updateScale(camera: THREE.Camera): void {
    if (!group.visible) return

    // Calculate distance from camera to gizmo
    const distance = camera.position.distanceTo(group.position)

    // For perspective camera, scale based on distance and FOV
    // For orthographic camera, scale based on zoom
    let scale: number
    if (camera instanceof THREE.PerspectiveCamera) {
      // Calculate the world size that corresponds to screenSize pixels at this distance
      const vFov = (camera.fov * Math.PI) / 180
      const worldHeight = 2 * Math.tan(vFov / 2) * distance
      // Assume a standard viewport height of ~800px for baseline
      scale = (worldHeight / 800) * screenSize * 0.5
    } else if (camera instanceof THREE.OrthographicCamera) {
      const viewHeight = camera.top - camera.bottom
      scale = (viewHeight / 800) * screenSize * 0.5
    } else {
      scale = 1
    }

    currentScale = scale
    group.scale.setScalar(scale)

    // Update arrow lengths to compensate for new scale
    updateArrowLengths()
  }

  function dispose(): void {
    // Dispose geometries and materials
    group.traverse((child) => {
      if (child instanceof THREE.Mesh) {
        child.geometry.dispose()
        if (child.material instanceof THREE.Material) {
          child.material.dispose()
        }
      }
    })
  }

  return {
    group,
    show,
    hide,
    setPositionAndBounds,
    getHoveredAxis,
    setHighlightedAxis,
    updateScale,
    dispose,
  }
}

export interface DragState {
  isDragging: boolean
  axis: 'x' | 'y' | 'z' | null
  isRightDrag: boolean
  startWorldPosition: THREE.Vector3
  startMousePosition: THREE.Vector2
  grabOffset: number // Offset along axis from selection center to grab point
  gridSize: number
}

export function createDragState(): DragState {
  return {
    isDragging: false,
    axis: null,
    isRightDrag: false,
    startWorldPosition: new THREE.Vector3(),
    startMousePosition: new THREE.Vector2(),
    grabOffset: 0,
    gridSize: 0.5,
  }
}

const GRID_SIZES = [0.1, 0.25, 0.5, 1.0, 2.0, 5.0]

export function cycleGridSize(currentSize: number, direction: number): number {
  const currentIndex = GRID_SIZES.indexOf(currentSize)
  if (currentIndex === -1) return GRID_SIZES[2]! // default to 0.5

  const newIndex = Math.max(0, Math.min(GRID_SIZES.length - 1, currentIndex + direction))
  return GRID_SIZES[newIndex]!
}

export function projectMouseToAxis(
  mouse: THREE.Vector2,
  camera: THREE.Camera,
  axisOrigin: THREE.Vector3,
  axis: 'x' | 'y' | 'z'
): THREE.Vector3 {
  // Create a ray from camera through mouse position
  const raycaster = new THREE.Raycaster()
  raycaster.setFromCamera(mouse, camera)

  // Create a plane perpendicular to the view direction that contains the axis line
  const cameraDirection = new THREE.Vector3()
  camera.getWorldDirection(cameraDirection)

  const axisDirection = new THREE.Vector3()
  if (axis === 'x') axisDirection.set(1, 0, 0)
  else if (axis === 'y') axisDirection.set(0, 1, 0)
  else axisDirection.set(0, 0, 1)

  // Find the plane that contains the axis and is most perpendicular to the camera
  // We use the plane defined by the axis and the camera direction
  const planeNormal = new THREE.Vector3()
    .crossVectors(axisDirection, cameraDirection)
    .cross(axisDirection)
    .normalize()

  // If plane normal is zero (looking straight down axis), use camera up
  if (planeNormal.lengthSq() < 0.001) {
    const cameraUp = new THREE.Vector3(0, 1, 0).applyQuaternion(camera.quaternion)
    planeNormal.crossVectors(axisDirection, cameraUp).cross(axisDirection).normalize()
  }

  const plane = new THREE.Plane().setFromNormalAndCoplanarPoint(planeNormal, axisOrigin)

  // Find intersection of ray with plane
  const intersection = new THREE.Vector3()
  raycaster.ray.intersectPlane(plane, intersection)

  if (!intersection) {
    return axisOrigin.clone()
  }

  // Project intersection onto the axis line
  const toIntersection = intersection.clone().sub(axisOrigin)
  const projectedDistance = toIntersection.dot(axisDirection)

  return axisOrigin.clone().add(axisDirection.multiplyScalar(projectedDistance))
}

export function getAxisDistance(
  mouse: THREE.Vector2,
  camera: THREE.Camera,
  axisOrigin: THREE.Vector3,
  axis: 'x' | 'y' | 'z'
): number {
  const projected = projectMouseToAxis(mouse, camera, axisOrigin, axis)
  const axisDirection = new THREE.Vector3()
  if (axis === 'x') axisDirection.set(1, 0, 0)
  else if (axis === 'y') axisDirection.set(0, 1, 0)
  else axisDirection.set(0, 0, 1)

  return projected.clone().sub(axisOrigin).dot(axisDirection)
}

export function snapToGrid(value: number, gridSize: number): number {
  return Math.round(value / gridSize) * gridSize
}