Dwinzo_Demo/app/src/modules/simulation/spatialUI/arm/useDraggableGLTF.ts

import { useRef, useState } from "react";
import * as THREE from "three";
import { ThreeEvent, useThree } from "@react-three/fiber";
import { useProductStore } from "../../../../store/simulation/useProductStore";
import {
    useSelectedEventData,
} from "../../../../store/simulation/useSimulationStore";
import { useProductContext } from "../../products/productContext";

type OnUpdateCallback = (object: THREE.Object3D) => void;

export default function useDraggableGLTF(onUpdate: OnUpdateCallback) {
    const { getEventByModelUuid } = useProductStore();
    const { selectedEventData } = useSelectedEventData();
    const { selectedProductStore } = useProductContext();
    const { selectedProduct } = selectedProductStore();
    const { camera, gl, controls } = useThree();
    const activeObjRef = useRef<THREE.Object3D | null>(null);
    const planeRef = useRef<THREE.Plane>(new THREE.Plane(new THREE.Vector3(0, 1, 0), 0));
    const offsetRef = useRef<THREE.Vector3>(new THREE.Vector3());
    const initialPositionRef = useRef<THREE.Vector3>(new THREE.Vector3());

    const raycaster = new THREE.Raycaster();
    const pointer = new THREE.Vector2();
    const [objectWorldPos, setObjectWorldPos] = useState(new THREE.Vector3());

    const handlePointerDown = (e: ThreeEvent<PointerEvent>) => {
        e.stopPropagation();

        let obj: THREE.Object3D | null = e.object;

        // Traverse up until we find modelUuid in userData
        while (obj && !obj.userData?.modelUuid) {
            obj = obj.parent;
        }

        if (!obj) return;

        // Disable orbit controls while dragging
        if (controls) (controls as any).enabled = false;

        activeObjRef.current = obj;
        initialPositionRef.current.copy(obj.position);

        // Get world position
        setObjectWorldPos(obj.getWorldPosition(objectWorldPos));

        // Set plane at the object's Y level
        planeRef.current.set(new THREE.Vector3(0, 1, 0), -objectWorldPos.y);

        // Convert pointer to NDC
        const rect = gl.domElement.getBoundingClientRect();
        pointer.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
        pointer.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;

        // Raycast to intersection
        raycaster.setFromCamera(pointer, camera);
        const intersection = new THREE.Vector3();
        raycaster.ray.intersectPlane(planeRef.current, intersection);

        // Calculate offset
        offsetRef.current.copy(objectWorldPos).sub(intersection);

        // Start listening for drag
        gl.domElement.addEventListener("pointermove", handlePointerMove);
        gl.domElement.addEventListener("pointerup", handlePointerUp);
    };

    const handlePointerMove = (e: PointerEvent) => {
        if (!activeObjRef.current) return;
        if (selectedEventData?.data.type === "roboticArm") {
            const selectedArmBot = getEventByModelUuid(
                selectedProduct.productUuid,
                selectedEventData.data.modelUuid
            );
            if (!selectedArmBot) return;
            // Check if Shift key is pressed
            const isShiftKeyPressed = e.shiftKey;

            // Get the mouse position relative to the canvas
            const rect = gl.domElement.getBoundingClientRect();
            pointer.x = ((e.clientX - rect.left) / rect.width) * 2 - 1;
            pointer.y = -((e.clientY - rect.top) / rect.height) * 2 + 1;

            // Update raycaster to point to the mouse position
            raycaster.setFromCamera(pointer, camera);

            // Create a vector to store intersection point
            const intersection = new THREE.Vector3();
            const intersects = raycaster.ray.intersectPlane(
                planeRef.current,
                intersection
            );
            if (!intersects) return;

            // Add offset for dragging
            intersection.add(offsetRef.current);

            // Get the parent's world matrix if exists
            const parent = activeObjRef.current.parent;
            const targetPosition = new THREE.Vector3();

            // OnPointerDown
            initialPositionRef.current.copy(objectWorldPos);

            // OnPointerMove
            if (isShiftKeyPressed) {
                const { x: initialX, y: initialY } = initialPositionRef.current;
                const { x: objectX, z: objectZ } = objectWorldPos;

                const deltaX = intersection.x - initialX;

                targetPosition.set(objectX, initialY + deltaX, objectZ);
            } else {
                // For free movement
                targetPosition.copy(intersection);
            }

            // CONSTRAIN MOVEMENT HERE:
            const centerX = selectedArmBot.position[0];
            const centerZ = selectedArmBot.position[2];
            const minDistance = 1.2;
            const maxDistance = 2;

            const delta = new THREE.Vector3(targetPosition.x - centerX, 0, targetPosition.z - centerZ);

            // Create quaternion from rotation
            const robotEuler = new THREE.Euler(selectedArmBot.rotation[0], selectedArmBot.rotation[1], selectedArmBot.rotation[2]);
            const robotQuaternion = new THREE.Quaternion().setFromEuler(robotEuler);

            // Inverse rotate
            const inverseQuaternion = robotQuaternion.clone().invert();
            delta.applyQuaternion(inverseQuaternion);

            // Angle in robot local space
            let relativeAngle = Math.atan2(delta.z, delta.x);
            let angleDeg = (relativeAngle * 180) / Math.PI;
            if (angleDeg < 0) {
                angleDeg += 360;
            }

            // Clamp angle
            if (angleDeg < 0 || angleDeg > 270) {
                const distanceTo90 = Math.abs(angleDeg - 0);
                const distanceTo270 = Math.abs(angleDeg - 270);
                if (distanceTo90 < distanceTo270) {
                    angleDeg = 0;
                } else {
                    return;
                }
                relativeAngle = (angleDeg * Math.PI) / 180;
            }

            // Distance clamp
            const distance = delta.length();
            const clampedDistance = Math.min(Math.max(distance, minDistance), maxDistance);

            // Calculate local target
            const finalLocal = new THREE.Vector3(
                Math.cos(relativeAngle) * clampedDistance,
                0,
                Math.sin(relativeAngle) * clampedDistance
            );

            // Rotate back to world space
            finalLocal.applyQuaternion(robotQuaternion);

            targetPosition.x = centerX + finalLocal.x;
            targetPosition.z = centerZ + finalLocal.z;


            // Clamp Y axis if needed
            targetPosition.y = Math.min(Math.max(targetPosition.y, 0.6), 1.9);

            // Convert to local if parent exists
            if (parent) {
                parent.worldToLocal(targetPosition);
            }

            // Update the object position
            activeObjRef.current.position.copy(targetPosition);
        }
    };

    const handlePointerUp = () => {
        if (controls) (controls as any).enabled = true;

        if (activeObjRef.current) {
            // Pass the updated position to the onUpdate callback to persist it
            onUpdate(activeObjRef.current);
        }

        gl.domElement.removeEventListener("pointermove", handlePointerMove);
        gl.domElement.removeEventListener("pointerup", handlePointerUp);

        activeObjRef.current = null;
    };

    return { handlePointerDown };
}