import { useCallback, useEffect, useRef, useState } from 'react';

export type DragLifecycleCustomState = Record<string, any>;

export interface DragLifecycleEventData<S extends DragLifecycleCustomState> {
  event: MouseEvent | TouchEvent;
  /**
   * Custom local state that is available in the callbacks
   */
  customState: S;
  /**
   * The x-coordinate when dragging started
   */
  startX: number;
  /**
   * The y-coordinate when dragging started
   */
  startY: number;
  /**
   * The x-coordinate of the most recent drag event
   */
  clientX: number;
  /**
   * The y-coordinate of the most recent drag event
   */
  clientY: number;
  /**
   * Difference in x-coordinate from when dragging start
   */
  deltaX: number;
  /**
   * Difference in y-coordinate from when dragging start
   */
  deltaY: number;
  /**
   * Difference in x-coordinate from the previous drag event
   */
  deltaXFromLast: number;
  /**
   * Difference in y-coordinate from the previous drag event
   */
  deltaYFromLast: number;
}

export type DragLifecycleCallback<S extends DragLifecycleCustomState = object> =
  (data: DragLifecycleEventData<S>) => void;

interface UseDragLifecycleOptions<S extends DragLifecycleCustomState> {
  customState?: S;
  /**
   * By default, the drag start will preventDefault to avoid text selection
   */
  shouldPreventDefault?:
    | boolean
    | ((
        e: MouseEvent | TouchEvent | React.MouseEvent | React.TouchEvent
      ) => boolean | undefined);
  /**
   * Determines whether the mousedown/touchstart event will start a drag lifecycle
   *
   * By default, this will filter out events that aren't touch or left mouse clicks.
   */
  shouldStartDrag?: (
    e: MouseEvent | TouchEvent | React.MouseEvent | React.TouchEvent
  ) => boolean | undefined;
  /**
   * Fires when dragging starts
   */
  onDragStart?: DragLifecycleCallback<S>;
  /**
   * Fires when moving during a drag
   */
  onDragMove?: DragLifecycleCallback<S>;
  /**
   * Fires when dragging ends
   */
  onDragEnd?: DragLifecycleCallback<S>;
  /**
   * Fires on any update to the drag position, including start, move, and end
   */
  onDragUpdate?: DragLifecycleCallback<S>;
}

type DragLifecycleReturn<S extends DragLifecycleCustomState> = [
  onDragStart: (
    e: React.MouseEvent | React.TouchEvent,
    context?: Partial<S>
  ) => void,
  cancel: () => void,
] & {
  /**
   * Event handler that invoked when touch/click starts a drag lifecycle.
   *
   * Adding and removing the touch/mouse move and end event handlers is handled
   * internally when dragging starts.
   */
  onDragStart: (
    e: React.MouseEvent | React.TouchEvent,
    context?: Partial<S>
  ) => void;
  /**
   * Forcibly interrupt a drag lifecycle
   */
  cancel: () => void;
};

export function getEventCoordinates(
  e: MouseEvent | TouchEvent
):
  | { clientX: number; clientY: number }
  | { clientX: undefined; clientY: undefined } {
  return 'touches' in e
    ? {
        clientX: e.touches[0]?.clientX,
        clientY: e.touches[0]?.clientY,
      }
    : {
        clientX: e.clientX,
        clientY: e.clientY,
      };
}

/**
 * Default behavior for allowing dragging is to accept touch events as well as
 * left clicks that aren't using modifier keys to emulate right clicks
 */
export function defaultShouldStartDrag(
  e: MouseEvent | TouchEvent | React.MouseEvent | React.TouchEvent
) {
  // Allow a touch event
  if ('touches' in e) return true;
  // Reject non-left-clicks
  return e.button === 0 && !e.ctrlKey;
}

/**
 * Creates a drag handler that can be used for click/touch-and-drag operations
 */
export default function useDragLifecycle<
  S extends DragLifecycleCustomState = DragLifecycleCustomState,
>(options?: UseDragLifecycleOptions<S>) {
  const {
    customState: initialCustomState = {} as S,
    shouldPreventDefault = true,
    shouldStartDrag = defaultShouldStartDrag,
    onDragStart,
    onDragMove,
    onDragEnd,
    onDragUpdate,
  } = options || {};

  const propEventHandlers = useRef({
    onDragStart,
    onDragMove,
    onDragEnd,
    onDragUpdate,
  });
  useEffect(() => {
    propEventHandlers.current.onDragStart = onDragStart;
    propEventHandlers.current.onDragMove = onDragMove;
    propEventHandlers.current.onDragEnd = onDragEnd;
    propEventHandlers.current.onDragUpdate = onDragUpdate;
  }, [onDragStart, onDragMove, onDragEnd, onDragUpdate]);

  const customStateRef = useRef(initialCustomState);

  const stateRef = useRef({
    startX: 0,
    startY: 0,
    clientX: 0,
    clientY: 0,
    deltaX: 0,
    deltaY: 0,
    deltaXFromLast: 0,
    deltaYFromLast: 0,
  });

  const [eventHandlers, setEventHandlers] = useState<{
    onMove?: (e: MouseEvent | TouchEvent) => void;
    onUp?: (e: MouseEvent | TouchEvent) => void;
  }>({});

  useEffect(() => {
    const { onMove, onUp } = eventHandlers;
    if (onMove) {
      document.addEventListener('mousemove', onMove);
      document.addEventListener('touchmove', onMove, { passive: false });
    }
    if (onUp) {
      document.addEventListener('mouseup', onUp);
      document.addEventListener('touchend', onUp, { passive: false });
    }
    return () => {
      if (onMove) {
        document.removeEventListener('mousemove', onMove);
        document.removeEventListener('touchmove', onMove);
      }
      if (onUp) {
        document.removeEventListener('mouseup', onUp);
        document.removeEventListener('touchend', onUp);
      }
    };
  }, [eventHandlers]);

  const handleCancel = useCallback(() => {
    setEventHandlers({
      onMove: undefined,
      onUp: undefined,
    });
  }, []);

  const updateDragStateFromEvent = useCallback(
    (e: MouseEvent | TouchEvent, reset = false): DragLifecycleEventData<S> => {
      const {
        clientX = stateRef.current.clientX,
        clientY = stateRef.current.clientY,
      } = getEventCoordinates(e);

      if (reset) {
        stateRef.current.deltaXFromLast = 0;
        stateRef.current.deltaYFromLast = 0;
      } else {
        stateRef.current.deltaXFromLast = reset
          ? 0
          : clientX - stateRef.current.clientX;
        stateRef.current.deltaYFromLast = reset
          ? 0
          : clientY - stateRef.current.clientY;
      }

      stateRef.current.clientX = clientX;
      stateRef.current.clientY = clientY;
      stateRef.current.deltaX = clientX - stateRef.current.startX;
      stateRef.current.deltaY = clientY - stateRef.current.startY;

      return {
        event: e,
        customState: customStateRef.current,
        ...stateRef.current,
      };
    },
    []
  );

  const handleDragStart = useCallback(
    (e: React.MouseEvent | React.TouchEvent, context?: Partial<S>) => {
      const allowDrag = shouldStartDrag(e);
      if (allowDrag === false) {
        setEventHandlers({
          onMove: undefined,
          onUp: undefined,
        });
      }

      if (!allowDrag) return;

      // Patch the context
      if (context) {
        Object.assign(customStateRef.current, context);
      }

      // Set the starting drag position
      const { clientX = 0, clientY = 0 } = getEventCoordinates(e.nativeEvent);

      stateRef.current.startX = clientX;
      stateRef.current.startY = clientY;

      const dragEventData = updateDragStateFromEvent(e.nativeEvent, true);

      propEventHandlers.current.onDragStart?.(dragEventData);
      propEventHandlers.current.onDragUpdate?.(dragEventData);

      if (e.isDefaultPrevented()) {
        // Event handlers can cancel drag
        setEventHandlers({
          onMove: undefined,
          onUp: undefined,
        });
      } else if (
        propEventHandlers.current.onDragMove ||
        propEventHandlers.current.onDragEnd ||
        propEventHandlers.current.onDragUpdate
      ) {
        if (
          typeof shouldPreventDefault === 'function'
            ? shouldPreventDefault(e)
            : shouldPreventDefault
        ) {
          e.preventDefault();
        }
        setEventHandlers({
          onMove(e: MouseEvent | TouchEvent) {
            const dragEventData = updateDragStateFromEvent(e);
            propEventHandlers.current.onDragMove?.(dragEventData);
            propEventHandlers.current.onDragUpdate?.(dragEventData);
            // Event handlers can interrupt drag
            if (e.defaultPrevented) {
              setEventHandlers({
                onMove: undefined,
                onUp: undefined,
              });
            }
          },
          onUp(e: MouseEvent | TouchEvent) {
            const dragEventData = updateDragStateFromEvent(e);
            propEventHandlers.current.onDragUpdate?.(dragEventData);
            propEventHandlers.current.onDragEnd?.(dragEventData);
            setEventHandlers({
              onMove: undefined,
              onUp: undefined,
            });
          },
        });
      }
    },
    [shouldPreventDefault, shouldStartDrag, updateDragStateFromEvent]
  );

  const result = [handleDragStart, handleCancel] as DragLifecycleReturn<S>;
  result.onDragStart = handleDragStart;
  result.cancel = handleCancel;
  return result;
}
