La-Fabrik/src/hooks/three/useTerrainHeight.ts

import { useMemo } from "react";
import { useGLTF } from "@react-three/drei";
import * as THREE from "three";
import { TERRAIN_MODEL_PATH } from "@/data/world/terrainConfig";
import {
  isInsideLaFabrikFootprint,
  LA_FABRIK_FLOOR_Y,
} from "@/data/world/laFabrikConfig";
import type { Vector3Tuple } from "@/types/three/three";
import { getMapNodesByName } from "@/utils/map/loadMapSceneData";

const RAYCAST_Y = 500;
const RAYCAST_FAR = 1000;
const DOWN = new THREE.Vector3(0, -1, 0);
const DEFAULT_TERRAIN_POSITION: Vector3Tuple = [0, 0, 0];
const DEFAULT_TERRAIN_ROTATION: Vector3Tuple = [0, 0, 0];
const DEFAULT_TERRAIN_SCALE: Vector3Tuple = [1, 1, 1];

interface TerrainHeightSampler {
  getHeight: (x: number, z: number) => number | null;
}

interface CachedTerrainHeightSampler {
  key: string;
  sampler: TerrainHeightSampler;
}

const terrainSamplerCache = new WeakMap<
  THREE.Object3D,
  CachedTerrainHeightSampler
>();

function createTerrainSamplerCacheKey(
  position: Vector3Tuple,
  rotation: Vector3Tuple,
  scale: Vector3Tuple,
): string {
  return `${position.join(",")}|${rotation.join(",")}|${scale.join(",")}`;
}

function createTerrainHeightSampler(
  scene: THREE.Object3D,
  position: Vector3Tuple,
  rotation: Vector3Tuple,
  scale: Vector3Tuple,
): TerrainHeightSampler {
  const meshes: THREE.Mesh[] = [];
  const terrainMatrix = new THREE.Matrix4().compose(
    new THREE.Vector3(...position),
    new THREE.Quaternion().setFromEuler(new THREE.Euler(...rotation)),
    new THREE.Vector3(...scale),
  );
  const inverseTerrainMatrix = terrainMatrix.clone().invert();
  const localOrigin = new THREE.Vector3();
  const localDirection = DOWN.clone().transformDirection(inverseTerrainMatrix);
  const hits: THREE.Intersection[] = [];
  const raycaster = new THREE.Raycaster(
    new THREE.Vector3(),
    DOWN,
    0,
    RAYCAST_FAR,
  );

  scene.updateMatrixWorld(true);
  scene.traverse((child) => {
    if (child instanceof THREE.Mesh) {
      meshes.push(child);
    }
  });

  return {
    getHeight: (x, z) => {
      if (isInsideLaFabrikFootprint(x, z, 0.6)) {
        return LA_FABRIK_FLOOR_Y;
      }

      localOrigin.set(x, RAYCAST_Y, z).applyMatrix4(inverseTerrainMatrix);
      raycaster.set(localOrigin, localDirection);
      hits.length = 0;
      raycaster.intersectObjects(meshes, false, hits);
      const hit = hits[0];
      return hit?.point.applyMatrix4(terrainMatrix).y ?? null;
    },
  };
}

export function useTerrainHeightSampler(): TerrainHeightSampler {
  const { scene } = useGLTF(TERRAIN_MODEL_PATH);
  const terrainNode = getMapNodesByName("terrain")[0];
  const position = terrainNode?.position ?? DEFAULT_TERRAIN_POSITION;
  const rotation = terrainNode?.rotation ?? DEFAULT_TERRAIN_ROTATION;
  const scale = terrainNode?.scale ?? DEFAULT_TERRAIN_SCALE;

  return useMemo(() => {
    const key = createTerrainSamplerCacheKey(position, rotation, scale);
    const cached = terrainSamplerCache.get(scene);

    if (cached?.key === key) {
      return cached.sampler;
    }

    const sampler = createTerrainHeightSampler(
      scene,
      position,
      rotation,
      scale,
    );
    terrainSamplerCache.set(scene, { key, sampler });
    return sampler;
  }, [position, rotation, scale, scene]);
}

export function useTerrainSnappedPosition(
  position: Vector3Tuple,
): Vector3Tuple {
  const terrainHeight = useTerrainHeightSampler();

  return useMemo(() => {
    const [x, y, z] = position;
    const height = terrainHeight.getHeight(x, z);
    return [x, height ?? y, z];
  }, [position, terrainHeight]);
}

export function getObjectBottomOffset(
  object: THREE.Object3D,
  scale: Vector3Tuple = [1, 1, 1],
): number {
  const bounds = new THREE.Box3().setFromObject(object);
  if (bounds.isEmpty()) return 0;

  return -bounds.min.y * scale[1];
}

export function normalizeMapScale(scale: Vector3Tuple): Vector3Tuple {
  const [x, y, z] = scale;
  const isUniform = Math.abs(x - y) < 0.001 && Math.abs(x - z) < 0.001;
  return isUniform ? scale : [x, x, x];
}