yslootahrobotics/src/components/RobotModel.tsx

'use client';

import { useRef, useMemo, useEffect, Suspense, useState, useCallback, Component, type ReactNode } from 'react';
import { useGLTF } from '@react-three/drei';
import { useFrame } from '@react-three/fiber';
import * as THREE from 'three';
import { useConfigStore } from '@/store/useConfigStore';
import { personaStore, usePersonaStore } from '@/store/usePersonaStore';

// Configure Draco decoder so compressed .glb files load correctly
useGLTF.setDecoderPath('/draco/');

const STATIC_ATTIRE_GLB: Record<string, string> = {
  'emarati-kandura': '/Kandoura.glb',
  'industrial-vest': '/Vest.glb',
  'business-suit': '/Suit.glb',
};

// Same chassis GLB for both — variants differ in software/licensing (EDU = open-source
// education / research build), not in geometry.
const BODY_GLB: Record<'basic' | 'edu', string> = {
  basic: '/Unitree_G1.glb',
  edu: '/Unitree_G1.glb',
};

// Attire models are loaded on-demand to avoid blocking the initial 50 MB robot load

/** Merge static map with any custom GLBs stored in the persona store */
function buildAttireGlbMap(personas: { id: string; modelPath?: string }[]): Record<string, string> {
  const dynamic: Record<string, string> = {};
  personas.forEach((p) => {
    if (p.modelPath && p.id !== 'none') dynamic[p.id] = p.modelPath;
  });
  return { ...STATIC_ATTIRE_GLB, ...dynamic }; // uploaded GLBs override static
}

interface ErrorBoundaryProps {
  children: ReactNode;
  onError: () => void;
  tag: string;
}
interface ErrorBoundaryState { hasError: boolean; }

class ModelErrorBoundary extends Component<ErrorBoundaryProps, ErrorBoundaryState> {
  constructor(props: ErrorBoundaryProps) {
    super(props);
    this.state = { hasError: false };
  }
  static getDerivedStateFromError(): ErrorBoundaryState {
    return { hasError: true };
  }
  componentDidCatch(error: Error) {
    console.warn(`[${this.props.tag}] Failed to load GLB, falling back:`, error.message);
    this.props.onError();
  }
  render() {
    if (this.state.hasError) return null;
    return this.props.children;
  }
}

function easeInOutCubic(t: number): number {
  return t < 0.5
    ? 4 * t * t * t
    : 1 - Math.pow(-2 * t + 2, 3) / 2;
}

function BaseBodyMesh({ glbPath, primaryColor, visible }: { glbPath: string; primaryColor: string; visible: boolean }) {
  useEffect(() => {
    return () => {
      useGLTF.clear(glbPath);
    };
  }, [glbPath]);

  const { scene } = useGLTF(glbPath);

  const processedScene = useMemo(() => {
    const clonedScene = scene.clone();

    clonedScene.traverse((child) => {
      if (child instanceof THREE.Mesh) {
        if (!child.material) {
          child.material = new THREE.MeshStandardMaterial({
            color: '#96a2b6',
            metalness: 0.8,
            roughness: 0.2,
          });
        }

        if (child.material instanceof THREE.MeshStandardMaterial) {
          child.material.envMapIntensity = 1;
          child.material.needsUpdate = true;
        }
      }
    });

    const box = new THREE.Box3().setFromObject(clonedScene);
    const center = box.getCenter(new THREE.Vector3());
    const size = box.getSize(new THREE.Vector3());

    const maxDim = Math.max(size.x, size.y, size.z);
    const scale = 2 / maxDim;
    clonedScene.scale.setScalar(scale);

    const offset = new THREE.Vector3(
      -center.x * scale,
      -center.y * scale + 0.5,
      -center.z * scale
    );
    clonedScene.position.copy(offset);

    return clonedScene;
  }, [scene]);

  useEffect(() => {
    processedScene.visible = visible;
  }, [visible, processedScene]);

  useEffect(() => {
    processedScene.traverse((child) => {
      if (child instanceof THREE.Mesh && child.material instanceof THREE.MeshStandardMaterial) {
        if (child.name.startsWith('Unitree_G1')) {
          child.material.color.set(primaryColor);
          child.material.needsUpdate = true;
        }
      }
    });
  }, [primaryColor, processedScene]);

  return <primitive object={processedScene} />;
}

function AttireModel({ glbPath, onLoaded }: { glbPath: string; onLoaded: () => void }) {
  // Clear stale useGLTF cache entries for paths that share the same base filename
  // (e.g. /models/robot-doctor.glb?v=1 replaced by ?v=2). This prevents Three.js
  // from serving the old binary when the version query string changes.
  useEffect(() => {
    return () => {
      useGLTF.clear(glbPath);
    };
  }, [glbPath]);

  const { scene } = useGLTF(glbPath);

  const processedAttire = useMemo(() => {
    const cloned = scene.clone();

    cloned.traverse((child) => {
      if (child instanceof THREE.Mesh && child.material instanceof THREE.MeshStandardMaterial) {
        child.material.envMapIntensity = 1;
        child.material.needsUpdate = true;
      }
    });

    const box = new THREE.Box3().setFromObject(cloned);
    const center = box.getCenter(new THREE.Vector3());
    const size = box.getSize(new THREE.Vector3());
    const maxDim = Math.max(size.x, size.y, size.z);
    const scale = 2 / maxDim;

    cloned.scale.setScalar(scale);
    cloned.position.set(
      -center.x * scale,
      -center.y * scale + 0.5,
      -center.z * scale
    );

    return cloned;
  }, [scene]);

  // Signal that the model is ready and rendered
  useEffect(() => {
    onLoaded();
  }, [onLoaded]);

  return <primitive object={processedAttire} />;
}

interface RobotModelProps {
  onError?: (error: Error) => void;
}

export function RobotModel({ onError: _onError }: RobotModelProps) {
  const groupRef = useRef<THREE.Group>(null);

  const spinningRef = useRef(false);
  const spinProgressRef = useRef(0);
  const swappedRef = useRef(false);
  const targetAttireRef = useRef<string>('none');

  const [displayedAttire, setDisplayedAttire] = useState('none');
  const [attireReady, setAttireReady] = useState(false);
  const previousAttireRef = useRef('none');

  const activeBody = useConfigStore((state) => state.activeBody);
  const bodyGlbPath = BODY_GLB[activeBody] ?? BODY_GLB.basic;

  const activeColors = useConfigStore((state) => state.activeColors);
  const activePersonaAttire = useConfigStore((state) => state.activePersonaAttire);
  // Subscribe to persona store so custom GLB paths are reactive
  const personas = usePersonaStore((s) => s.personas);

  // Detect attire change and trigger spin
  useEffect(() => {
    if (activePersonaAttire !== previousAttireRef.current) {
      targetAttireRef.current = activePersonaAttire;
      spinningRef.current = true;
      spinProgressRef.current = 0;
      swappedRef.current = false;
      setAttireReady(false);
      previousAttireRef.current = activePersonaAttire;
    }
  }, [activePersonaAttire]);

  // Spin animation loop
  useFrame((_, delta) => {
    if (!spinningRef.current || !groupRef.current) return;

    const spinSpeed = 2.5;
    spinProgressRef.current += delta * spinSpeed;

    const progress = Math.min(spinProgressRef.current, 1);
    const easedRotation = easeInOutCubic(progress) * Math.PI * 2;
    groupRef.current.rotation.y = easedRotation;

    if (progress >= 0.5 && !swappedRef.current) {
      setDisplayedAttire(targetAttireRef.current);
      swappedRef.current = true;
    }

    if (progress >= 1) {
      groupRef.current.rotation.y = 0;
      spinningRef.current = false;
      spinProgressRef.current = 0;
    }
  });

  // Hide base robot only when attire is selected AND the attire GLB has loaded
  // Show base robot when 'none' is selected or attire is still loading
  const showBase = displayedAttire === 'none' || !attireReady;

  const attireGlbPath = buildAttireGlbMap(personas)[displayedAttire] || null;

  const handleAttireLoaded = useCallback(() => {
    setAttireReady(true);
  }, []);

  return (
    <group ref={groupRef}>
      <Suspense fallback={null}>
        <BaseBodyMesh glbPath={bodyGlbPath} primaryColor={activeColors.primary} visible={showBase} />
      </Suspense>

      {attireGlbPath && (
        <ModelErrorBoundary key={attireGlbPath} tag="AttireModel" onError={() => { setDisplayedAttire('none'); setAttireReady(false); }}>
          <Suspense fallback={null}>
            <AttireModel
              glbPath={attireGlbPath}
              onLoaded={handleAttireLoaded}
            />
          </Suspense>
        </ModelErrorBoundary>
      )}
    </group>
  );
}

useGLTF.preload('/Unitree_G1.glb');
// Preload uploaded attire in the background so they're ready before the user clicks
useGLTF.preload('/models/robot-doctor.glb');
useGLTF.preload('/models/security-guard.glb');