import {
  Float32BufferAttribute,
  Group,
  Mesh,
  type Object3D,
} from "three";

export interface PreparedCollisionTrimesh {
  readonly id: string;
  readonly vertices: ArrayLike<number>;
  readonly indices: ArrayLike<number>;
}

export function prepareCollision(root: Object3D): Group {
  root.updateWorldMatrix(true, true);
  const inverseRootMatrix = root.matrixWorld.clone().invert();
  const prepared = new Group();
  prepared.name = `${root.name || "collision"}-prepared`;
  root.traverse((node) => {
    const mesh = node as Mesh;
    if (!mesh.isMesh) return;
    const geometry = mesh.geometry.clone();
    const position = geometry.getAttribute("position");
    if (position.normalized || !(position.array instanceof Float32Array)) {
      const dequantized = new Float32Array(position.count * position.itemSize);
      for (let vertex = 0; vertex < position.count; vertex += 1) {
        const offset = vertex * position.itemSize;
        dequantized[offset] = position.getX(vertex);
        if (position.itemSize > 1) dequantized[offset + 1] = position.getY(vertex);
        if (position.itemSize > 2) dequantized[offset + 2] = position.getZ(vertex);
        if (position.itemSize > 3) dequantized[offset + 3] = position.getW(vertex);
      }
      geometry.setAttribute(
        "position",
        new Float32BufferAttribute(dequantized, position.itemSize),
      );
    }
    geometry.applyMatrix4(inverseRootMatrix.clone().multiply(mesh.matrixWorld));
    const preparedMesh = new Mesh(geometry, mesh.material);
    preparedMesh.name = mesh.name;
    preparedMesh.visible = false;
    prepared.add(preparedMesh);
  });
  return prepared;
}

export function collisionTrimeshes(root: Object3D): readonly PreparedCollisionTrimesh[] {
  const result: PreparedCollisionTrimesh[] = [];
  root.traverse((node) => {
    const mesh = node as Mesh;
    if (!mesh.isMesh) return;
    const position = mesh.geometry.getAttribute("position");
    const index = mesh.geometry.index;
    result.push({
      id: mesh.uuid,
      vertices: position.array,
      indices: index?.array
        ?? Uint32Array.from({ length: position.count }, (_, vertex) => vertex),
    });
  });
  return result;
}