import { type ReactNode, useCallback, useRef, useState } from 'react';

import {
  type CollisionDetection,
  closestCenter,
  DndContext,
  type DragEndEvent,
  type DragOverEvent,
  DragOverlay,
  type DragStartEvent,
  getFirstCollision,
  KeyboardSensor,
  type Modifiers,
  PointerSensor,
  pointerWithin,
  rectIntersection,
  type UniqueIdentifier,
  useSensor,
  useSensors,
} from '@dnd-kit/core';
import {
  restrictToHorizontalAxis,
  restrictToParentElement,
  restrictToVerticalAxis,
} from '@dnd-kit/modifiers';
import { sortableKeyboardCoordinates } from '@dnd-kit/sortable';
import Portal from '@mui/material/Portal';

import {
  CUSTOM_DROP_ANIMATION_CONFIG,
  DEFAULT_AUTO_SCROLL_CONFIG,
  DEFAULT_DRAG_ACTIVATION_DISTANCE,
  ROOT_CONTAINER_ID,
} from '../../constants';
import { SortableListContext } from '../../context';
import { type RootProps } from '../../types';
import Container from '../Container';
import OverlayContent from '../OverlayContent';

const Root = ({
  children,
  onDragEnd,
  onDragOver,
  direction = 'vertical',
  dragByHandler = false,
  restrictToAncestor = false,
  dragActivationDistance = DEFAULT_DRAG_ACTIVATION_DISTANCE,
  hasDragOverlay = false,
  sx = {},
  isDraggable,
}: RootProps) => {
  const [activeId, setActiveId] = useState<UniqueIdentifier | null>(null);
  const [containers, setContainers] = useState<Map<string, UniqueIdentifier[]>>(
    new Map(),
  );
  const [items, setItems] = useState<Map<UniqueIdentifier, ReactNode>>(
    new Map(),
  );
  const itemRefs = useRef<Map<UniqueIdentifier, HTMLElement>>(new Map());

  const sensors = useSensors(
    useSensor(PointerSensor, {
      activationConstraint: { distance: dragActivationDistance },
    }),
    useSensor(KeyboardSensor, {
      coordinateGetter: sortableKeyboardCoordinates,
    }),
  );

  const lastOverIdRef = useRef<UniqueIdentifier | null>(null);
  // Custom collision detection strategy for multiple containers
  const customCollisionStrategy: CollisionDetection = useCallback(
    (args: Parameters<CollisionDetection>[0]) => {
      const pointerCollisions = pointerWithin(args);
      const intersections =
        pointerCollisions.length > 0
          ? pointerCollisions
          : rectIntersection(args);

      const overId = getFirstCollision(intersections, 'id');

      if (overId != null) {
        lastOverIdRef.current = overId;
        return [{ id: overId }];
      }

      if (lastOverIdRef.current) {
        return [{ id: lastOverIdRef.current }];
      }

      return [];
    },
    [],
  );

  const collisionDetectionStrategy: CollisionDetection = hasDragOverlay
    ? closestCenter
    : customCollisionStrategy;

  const registerContainer = useCallback(
    (id: string, itemIds: UniqueIdentifier[]) => {
      setContainers(prev => {
        const newMap = new Map(prev);
        newMap.set(id, itemIds);
        return newMap;
      });
    },
    [],
  );

  const unregisterContainer = useCallback((id: string) => {
    setContainers(prev => {
      const newMap = new Map(prev);
      newMap.delete(id);
      return newMap;
    });
  }, []);

  const getContainerItems = useCallback(
    (id: string) => containers.get(id) || [],
    [containers],
  );

  const registerItem = useCallback(
    (
      id: UniqueIdentifier,
      content: ReactNode,
      element?: HTMLElement | null,
    ) => {
      setItems(prev => {
        const newMap = new Map(prev);
        newMap.set(id, content);
        return newMap;
      });
      if (element) {
        itemRefs.current.set(id, element);
      }
    },
    [],
  );

  const unregisterItem = useCallback((id: UniqueIdentifier) => {
    setItems(prev => {
      const newMap = new Map(prev);
      newMap.delete(id);
      return newMap;
    });
    itemRefs.current.delete(id);
  }, []);

  const findContainer = useCallback(
    (id: UniqueIdentifier) => {
      for (const [containerId, containerItems] of containers.entries()) {
        if (containerItems.includes(id)) {
          return containerId;
        }
      }
      return undefined;
    },
    [containers],
  );

  const handleDragStart = useCallback((event: DragStartEvent) => {
    setActiveId(event.active.id);
  }, []);

  const handleDragOver = useCallback(
    (event: DragOverEvent) => {
      const { active, over } = event;

      if (!over) return;

      const activeContainer = findContainer(active.id);
      const overContainer = findContainer(over.id);

      if (onDragOver) {
        onDragOver({
          active: active.id,
          over: over.id,
          activeContainer,
          overContainer,
        });
      }
    },
    [findContainer, onDragOver],
  );

  const handleDragEnd = useCallback(
    (event: DragEndEvent) => {
      const { active, over } = event;

      if (!over) {
        setActiveId(null);
        return;
      }

      const activeContainer = findContainer(active.id);
      const overContainer = findContainer(over.id);

      if (onDragEnd) {
        onDragEnd({
          active: active.id,
          over: over.id,
          activeContainer,
          overContainer,
        });
      }

      setActiveId(null);
    },
    [findContainer, onDragEnd],
  );

  const handleDragCancel = useCallback(() => {
    setActiveId(null);
  }, []);

  const activeElement = activeId ? items.get(activeId) : null;

  const modifiers: Modifiers = [
    direction === 'vertical'
      ? restrictToVerticalAxis
      : restrictToHorizontalAxis,
    ...(restrictToAncestor && !hasDragOverlay ? [restrictToParentElement] : []),
  ];

  return (
    <SortableListContext.Provider
      value={{
        direction,
        dragByHandler,
        hasDragOverlay,
        activeId,
        isDraggable,
        registerContainer,
        unregisterContainer,
        getContainerItems,
        registerItem,
        unregisterItem,
      }}
    >
      <DndContext
        sensors={sensors}
        collisionDetection={collisionDetectionStrategy}
        autoScroll={DEFAULT_AUTO_SCROLL_CONFIG}
        onDragStart={handleDragStart}
        onDragOver={handleDragOver}
        onDragEnd={handleDragEnd}
        onDragCancel={handleDragCancel}
        modifiers={modifiers}
      >
        <Container
          id={ROOT_CONTAINER_ID}
          sx={sx}
          overlay={
            hasDragOverlay && activeId && activeElement ? (
              <Portal>
                <DragOverlay dropAnimation={CUSTOM_DROP_ANIMATION_CONFIG}>
                  <OverlayContent id={activeId}>{activeElement}</OverlayContent>
                </DragOverlay>
              </Portal>
            ) : undefined
          }
        >
          {children}
        </Container>
      </DndContext>
    </SortableListContext.Provider>
  );
};

Root.displayName = 'SortableList.Root';

export default Root;
