healer-man/src/avatar/equipmentAttachments.tsx

import { useGLTF } from "@react-three/drei";
import { resolveContentUrl } from "../content/contentManager";
import { createPortal } from "@react-three/fiber";
import {
  Suspense,
  useEffect,
  useMemo,
  type ReactNode,
} from "react";
import {
  DoubleSide,
  Mesh,
  Object3D,
  type Material,
} from "three";
import { clone as cloneSkeleton } from "three/examples/jsm/utils/SkeletonUtils.js";
import type { AvatarManifestEntry } from "./characterManifest";
import type {
  EquipmentAttachmentName,
  EquipmentModelVisual,
} from "./equipmentAppearance";
import {
  resolveAvatarAttachmentPoint,
  type AvatarAttachmentTransform,
} from "./attachmentPoints";
import { AvatarErrorBoundary } from "./AvatarErrorBoundary";
import {
  avatarWeaponSocketRotation,
  findAvatarHandBone,
  findAvatarShieldBone,
  StarterWeaponVisual,
  type AvatarHand,
  type StarterWeaponRotation,
} from "./starterWeapons";

const EQUIPMENT_ATTACHMENT_NAMES: readonly EquipmentAttachmentName[] = Object.freeze([
  "handLeft",
  "handRight",
  "shield",
  "shoulderLeft",
  "shoulderRight",
]);
const NO_ROTATION: StarterWeaponRotation = Object.freeze([0, 0, 0]);

export type EquipmentAttachmentTransforms = Readonly<
  Partial<Record<EquipmentAttachmentName, AvatarAttachmentTransform>>
>;

export function equipmentAttachmentHand(
  attachment: EquipmentAttachmentName,
): AvatarHand | null {
  if (attachment === "handLeft" || attachment === "shield") return "left";
  if (attachment === "handRight") return "right";
  return null;
}

function cloneEquipmentScene(source: Object3D, name: string): Object3D {
  const scene = cloneSkeleton(source);
  scene.name = `equipped-${name}`;
  scene.traverse((node) => {
    const mesh = node as Mesh;
    if (!mesh.isMesh) return;
    const sourceMaterials = Array.isArray(mesh.material) ? mesh.material : [mesh.material];
    const materials = sourceMaterials.map((sourceMaterial) => {
      const material = sourceMaterial.clone();
      material.side = DoubleSide;
      material.needsUpdate = true;
      return material;
    });
    mesh.material = Array.isArray(mesh.material) ? materials : materials[0];
    mesh.castShadow = false;
    mesh.receiveShadow = false;
  });
  return scene;
}

function fallbackTarget(
  root: Object3D,
  attachment: EquipmentAttachmentName,
): Object3D | null {
  if (attachment === "handLeft" || attachment === "shield") {
    return attachment === "shield"
      ? findAvatarShieldBone(root)
      : findAvatarHandBone(root, "left");
  }
  if (attachment === "handRight") return findAvatarHandBone(root, "right");
  return root.getObjectByName(
    attachment === "shoulderLeft" ? "bone_ShoulderL" : "bone_ShoulderR",
  ) ?? null;
}

/** Snapshots all authored attachment offsets once, before animation begins. */
export function resolveEquipmentAttachmentTransforms(
  root: Object3D,
  entry: AvatarManifestEntry,
): EquipmentAttachmentTransforms {
  const transforms: Partial<Record<EquipmentAttachmentName, AvatarAttachmentTransform>> = {};
  for (const attachment of EQUIPMENT_ATTACHMENT_NAMES) {
    const transform = resolveAvatarAttachmentPoint(
      root,
      entry.attachmentPoints?.[attachment],
      fallbackTarget(root, attachment),
      {
        useTargetOrigin: attachment === "handLeft"
          || attachment === "handRight"
          || attachment === "shield",
      },
    );
    if (transform) transforms[attachment] = transform;
  }
  return transforms;
}

function AttachedEquipmentModel({
  entry,
  itemModel,
  transform,
}: {
  readonly entry: AvatarManifestEntry;
  readonly itemModel: EquipmentModelVisual;
  readonly transform: AvatarAttachmentTransform;
}): ReactNode {
  const gltf = useGLTF(resolveContentUrl(itemModel.url));
  const scene = useMemo(
    () => cloneEquipmentScene(gltf.scene, `${itemModel.itemId}-${itemModel.attachment}`),
    [gltf.scene, itemModel.attachment, itemModel.itemId],
  );
  useEffect(() => () => {
    const materials = new Set<Material>();
    scene.traverse((node) => {
      const mesh = node as Mesh;
      if (!mesh.isMesh) return;
      const meshMaterials = Array.isArray(mesh.material) ? mesh.material : [mesh.material];
      meshMaterials.forEach((material) => materials.add(material));
    });
    materials.forEach((material) => material.dispose());
  }, [scene]);

  const socketRotation = itemModel.fallbackKind && itemModel.weaponStance
    ? avatarWeaponSocketRotation(entry, itemModel.fallbackKind, itemModel.weaponStance)
    : NO_ROTATION;

  return createPortal(
    <group
      position={transform.position}
      rotation={socketRotation}
      scale={transform.scale}
    >
      <primitive object={scene} />
    </group>,
    transform.target,
  );
}

function EquipmentModelFallback({
  entry,
  fallbackScale,
  itemModel,
  transform,
}: {
  readonly entry: AvatarManifestEntry;
  readonly fallbackScale: number;
  readonly itemModel: EquipmentModelVisual;
  readonly transform: AvatarAttachmentTransform;
}): ReactNode {
  const hand = equipmentAttachmentHand(itemModel.attachment);
  if (!hand || !itemModel.fallbackKind || !itemModel.weaponStance) return null;
  const socketRotation = avatarWeaponSocketRotation(
    entry,
    itemModel.fallbackKind,
    itemModel.weaponStance,
  );
  return createPortal(
    <group
      position={transform.position}
      rotation={socketRotation}
      scale={transform.scale}
    >
      <StarterWeaponVisual
        hand={hand}
        kind={itemModel.fallbackKind}
        name={`fallback-equipment-${itemModel.itemId}-${itemModel.attachment}`}
        scale={fallbackScale}
      />
    </group>,
    transform.target,
  );
}

function ResolvedEquipmentAttachment({
  entry,
  fallbackScale,
  itemModel,
  transform,
}: {
  readonly entry: AvatarManifestEntry;
  readonly fallbackScale: number;
  readonly itemModel: EquipmentModelVisual;
  readonly transform: AvatarAttachmentTransform;
}): ReactNode {
  const fallback = (
    <EquipmentModelFallback
      entry={entry}
      fallbackScale={fallbackScale}
      itemModel={itemModel}
      transform={transform}
    />
  );
  return (
    <AvatarErrorBoundary
      fallback={fallback}
      resetKey={`${itemModel.itemId}:${itemModel.attachment}:${itemModel.url}`}
    >
      <Suspense fallback={fallback}>
        <AttachedEquipmentModel entry={entry} itemModel={itemModel} transform={transform} />
      </Suspense>
    </AvatarErrorBoundary>
  );
}

export function EquipmentAttachments({
  entry,
  fallbackScale,
  models,
  transforms,
}: {
  readonly entry: AvatarManifestEntry;
  readonly fallbackScale: number;
  readonly models: readonly EquipmentModelVisual[];
  readonly transforms: EquipmentAttachmentTransforms;
}): ReactNode {
  return models.map((itemModel) => {
    const transform = transforms[itemModel.attachment];
    return transform ? (
      <ResolvedEquipmentAttachment
        key={`${itemModel.itemId}:${itemModel.attachment}`}
        entry={entry}
        fallbackScale={fallbackScale}
        itemModel={itemModel}
        transform={transform}
      />
    ) : null;
  });
}