webeditor3d/src/geometry/box-face-uvs.ts

import { BoxGeometry } from "three";

import type { Vec2, Vec3 } from "../core/vector";
import {
  BOX_FACE_IDS,
  createDefaultFaceUvState,
  type Brush,
  type BoxFaceId,
  type FaceUvState,
  type WhiteboxFaceId
} from "../document/brushes";
import { getBrushFaceBasis, getBrushLocalVertexPosition } from "./whitebox-brush";
import { getBrushFaceVertexIds } from "./whitebox-topology";

interface BoxBrushUvProjectionSource {
  size: Vec3;
  faces: Record<BoxFaceId, { uv: FaceUvState }>;
}

function computeGenericBrushFaceSize(brush: Brush, faceId: WhiteboxFaceId): Vec2 {
  const basis = getBrushFaceBasis(brush, faceId);
  const projectedVertices = getBrushFaceVertexIds(brush, faceId).map(
    (vertexId) => {
      const vertex = getBrushLocalVertexPosition(brush, vertexId);
      const relative = {
        x: vertex.x - basis.origin.x,
        y: vertex.y - basis.origin.y,
        z: vertex.z - basis.origin.z
      };

      return {
        x:
          relative.x * basis.uAxis.x +
          relative.y * basis.uAxis.y +
          relative.z * basis.uAxis.z,
        y:
          relative.x * basis.vAxis.x +
          relative.y * basis.vAxis.y +
          relative.z * basis.vAxis.z
      };
    }
  );
  const firstVertex = projectedVertices[0];
  const min = { ...firstVertex };
  const max = { ...firstVertex };

  for (const projectedVertex of projectedVertices.slice(1)) {
    min.x = Math.min(min.x, projectedVertex.x);
    min.y = Math.min(min.y, projectedVertex.y);
    max.x = Math.max(max.x, projectedVertex.x);
    max.y = Math.max(max.y, projectedVertex.y);
  }

  return {
    x: max.x - min.x,
    y: max.y - min.y
  };
}

export function getBoxBrushFaceSize(
  brush: BoxBrushUvProjectionSource | Brush,
  faceId: BoxFaceId | WhiteboxFaceId
): Vec2 {
  if ("kind" in brush) {
    if (brush.kind !== "box") {
      return computeGenericBrushFaceSize(brush, faceId);
    }
  }

  switch (faceId) {
    case "posX":
    case "negX":
      return {
        x: brush.size.z,
        y: brush.size.y
      };
    case "posY":
    case "negY":
      return {
        x: brush.size.x,
        y: brush.size.z
      };
    case "posZ":
    case "negZ":
      return {
        x: brush.size.x,
        y: brush.size.y
      };
  }

  if ("kind" in brush) {
    return computeGenericBrushFaceSize(brush, faceId);
  }

  throw new Error(`Unsupported box face id ${faceId}.`);
}

export function createFitToFaceBoxBrushFaceUvState(
  brush: Brush,
  faceId: WhiteboxFaceId
): FaceUvState {
  const faceSize = getBoxBrushFaceSize(brush, faceId);

  return {
    ...createDefaultFaceUvState(),
    scale: {
      x: 1 / faceSize.x,
      y: 1 / faceSize.y
    }
  };
}

export function createFitToMaterialTileBoxBrushFaceUvState(
  brush: Brush,
  faceId: WhiteboxFaceId,
  tileSize: Vec2
): FaceUvState {
  const faceSize = getBoxBrushFaceSize(brush, faceId);

  return {
    ...createDefaultFaceUvState(),
    scale: {
      x: tileSize.x / faceSize.x,
      y: tileSize.y / faceSize.y
    }
  };
}

export function projectBoxFaceVertexToUv(vertexPosition: Vec3, brush: BoxBrushUvProjectionSource, faceId: BoxFaceId): Vec2 {
  const halfSize = {
    x: brush.size.x * 0.5,
    y: brush.size.y * 0.5,
    z: brush.size.z * 0.5
  };

  switch (faceId) {
    case "posX":
      return {
        x: halfSize.z - vertexPosition.z,
        y: vertexPosition.y + halfSize.y
      };
    case "negX":
      return {
        x: vertexPosition.z + halfSize.z,
        y: vertexPosition.y + halfSize.y
      };
    case "posY":
      return {
        x: vertexPosition.x + halfSize.x,
        y: halfSize.z - vertexPosition.z
      };
    case "negY":
      return {
        x: vertexPosition.x + halfSize.x,
        y: vertexPosition.z + halfSize.z
      };
    case "posZ":
      return {
        x: vertexPosition.x + halfSize.x,
        y: vertexPosition.y + halfSize.y
      };
    case "negZ":
      return {
        x: halfSize.x - vertexPosition.x,
        y: vertexPosition.y + halfSize.y
      };
  }
}

export function transformProjectedFaceUv(baseUv: Vec2, faceSize: Vec2, uvState: FaceUvState): Vec2 {
  let u = (baseUv.x - faceSize.x * 0.5) * uvState.scale.x;
  let v = (baseUv.y - faceSize.y * 0.5) * uvState.scale.y;

  if (uvState.flipU) {
    u *= -1;
  }

  if (uvState.flipV) {
    v *= -1;
  }

  switch (uvState.rotationQuarterTurns) {
    case 1: {
      const nextU = -v;
      v = u;
      u = nextU;
      break;
    }
    case 2:
      u *= -1;
      v *= -1;
      break;
    case 3: {
      const nextU = v;
      v = -u;
      u = nextU;
      break;
    }
  }

  return {
    x: u + faceSize.x * 0.5 * uvState.scale.x + uvState.offset.x,
    y: v + faceSize.y * 0.5 * uvState.scale.y + uvState.offset.y
  };
}

export function applyBoxBrushFaceUvsToGeometry(geometry: BoxGeometry, brush: BoxBrushUvProjectionSource): void {
  const positionAttribute = geometry.getAttribute("position");
  const uvAttribute = geometry.getAttribute("uv");
  const indexAttribute = geometry.getIndex();

  if (indexAttribute === null) {
    throw new Error("BoxGeometry is expected to be indexed for face UV projection.");
  }

  // BoxGeometry groups follow the same px, nx, py, ny, pz, nz order as the canonical face ids.
  for (const [materialIndex, faceId] of BOX_FACE_IDS.entries()) {
    const group = geometry.groups.find((candidate) => candidate.materialIndex === materialIndex);

    if (group === undefined) {
      continue;
    }

    const faceSize = getBoxBrushFaceSize(brush, faceId);
    const vertexIndices = new Set<number>();

    for (let indexOffset = group.start; indexOffset < group.start + group.count; indexOffset += 1) {
      vertexIndices.add(indexAttribute.getX(indexOffset));
    }

    for (const vertexIndex of vertexIndices) {
      const localVertexPosition = {
        x: positionAttribute.getX(vertexIndex),
        y: positionAttribute.getY(vertexIndex),
        z: positionAttribute.getZ(vertexIndex)
      };
      const projectedUv = projectBoxFaceVertexToUv(localVertexPosition, brush, faceId);
      const transformedUv = transformProjectedFaceUv(projectedUv, faceSize, brush.faces[faceId].uv);

      uvAttribute.setXY(vertexIndex, transformedUv.x, transformedUv.y);
    }
  }

  uvAttribute.needsUpdate = true;
}