import { useRef } from "react";
import * as THREE from "three";
import { ThreeEvent, useThree } from "@react-three/fiber";

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

export default function useDraggableGLTF(onUpdate: OnUpdateCallback) {
  const { camera, gl, controls, scene } = 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 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
    const objectWorldPos = new THREE.Vector3();
    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;
  
    // 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);
    console.log('intersection: ', intersection);
  
    // Get the parent's world matrix if exists
    const parent = activeObjRef.current.parent;
    const targetPosition = new THREE.Vector3();
    
    if (isShiftKeyPressed) {
        console.log('isShiftKeyPressed: ', isShiftKeyPressed);
      // For Y-axis only movement, maintain original X and Z
      console.log('initialPositionRef: ', initialPositionRef);
      console.log('intersection.y: ', intersection);
      targetPosition.set(
        initialPositionRef.current.x,
        intersection.y,
        initialPositionRef.current.z
      );
    } else {
      // For free movement
      targetPosition.copy(intersection);
    }
  
    // Convert world position to local if object is nested inside a parent
    if (parent) {
      parent.worldToLocal(targetPosition);
    }
  
    // Update 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 };
}