import type React from 'react';
import {
  createContext,
  type ReactNode,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';

import { WeekMonthPeriod } from '@material-hu/components/composed-components/WeekMonthSelector/types';

import {
  type CopiedShiftCell,
  type ShiftAssignCellVC,
  type ShiftCalendarUser,
} from 'src/types/shifts';
import { getWeekRange } from 'src/utils/date';

import {
  checkCellSelectionPromotions,
  filterCellsByCollaborators,
  filterCellsByDate,
  filterShiftsByCells,
  filterShiftsByCollaborator,
  filterShiftsByDate,
  filterShiftsByIds,
} from '../utils';

export type ShiftSelection = {
  shiftId: number;
  userId: number;
  date: string;
};

type ShiftsSelectionState = {
  // Modes
  isAdvancedMode: boolean;
  isAdvancedModeManual: boolean; //Advanced mode was manually toggled
  allRowsSelected: boolean;
  viewMode: WeekMonthPeriod;

  startDate: Date;
  endDate: Date;
  selectedDates: Set<string>;
  selectedCells: ShiftAssignCellVC[];
  excludeCells: ShiftAssignCellVC[];
  selectedShifts: Map<number, ShiftSelection>;
  excludeShifts: Map<number, ShiftSelection>;

  // Copy/Paste
  copiedCells: CopiedShiftCell[];
  pastingCells: Set<string>;
};

// Context
const ShiftsSelectionContext = createContext<{
  state: ShiftsSelectionState;
  setState: React.Dispatch<React.SetStateAction<ShiftsSelectionState>>;
  closeOpenOverlays: () => void;
  registerCloseOverlay: (closeFn: () => void) => () => void;
} | null>(null);

// Provider
type ShiftsSelectionProviderProps = {
  children: ReactNode;
};

type ClearAllSelectionsParams = {
  keepDates?: boolean;
  keepCopiedCells?: boolean;
  keepAdvancedMode?: boolean;
  keepSelectedShifts?: boolean;
};

export const ShiftsSelectionProvider: React.FC<
  ShiftsSelectionProviderProps
> = ({ children }) => {
  const [state, setState] = useState<ShiftsSelectionState>(() => {
    const today = new Date();
    const weekRange = getWeekRange(today);
    return {
      isAdvancedMode: false,
      isAdvancedModeManual: false,
      viewMode: WeekMonthPeriod.WEEKLY,
      startDate: weekRange.startDate,
      endDate: weekRange.endDate,
      selectedDates: new Set(),
      selectedCells: [],
      excludeCells: [],
      allRowsSelected: false,
      selectedShifts: new Map(),
      excludeShifts: new Map(),
      copiedCells: [],
      pastingCells: new Set(),
    };
  });

  const overlayCloseCallbacksRef = useRef<Set<() => void>>(new Set());

  const closeOpenOverlays = useCallback(() => {
    overlayCloseCallbacksRef.current.forEach(closeFn => {
      closeFn();
    });
  }, []);

  const registerCloseOverlay = useCallback((closeFn: () => void) => {
    overlayCloseCallbacksRef.current.add(closeFn);
    return () => {
      overlayCloseCallbacksRef.current.delete(closeFn);
    };
  }, []);

  // Memoize the context value to prevent unnecessary re-renders
  const contextValue = useMemo(
    () => ({ state, setState, closeOpenOverlays, registerCloseOverlay }),
    [state, setState, closeOpenOverlays, registerCloseOverlay],
  );

  return (
    <ShiftsSelectionContext.Provider value={contextValue}>
      {children}
    </ShiftsSelectionContext.Provider>
  );
};

// Context hook
export const useShiftsSelection = () => {
  const context = useContext(ShiftsSelectionContext);
  if (!context) {
    throw new Error(
      'useShiftsSelection must be used within a ShiftsSelectionProvider',
    );
  }
  return context;
};

const calculateAdvancedModeState = (
  hasSelections: boolean,
  currentState: ShiftsSelectionState,
) => {
  // Advanced mode should be enabled if:
  // 1. It was manually toggled on, OR
  // 2. There are active selections
  const isAdvancedMode = currentState.isAdvancedModeManual || hasSelections;

  // If there are no selections, we also clear any manual toggle so that
  // hitting the last deselect returns the UI to normal mode.  Manual
  // toggles only persist while there is some state to work with.
  const isAdvancedModeManual = hasSelections
    ? currentState.isAdvancedModeManual
    : false;

  // Clear shift selections when enabling advanced mode (transition from false to true)
  const shouldClearShiftSelections =
    !currentState.isAdvancedMode && isAdvancedMode;

  return {
    isAdvancedMode,
    isAdvancedModeManual,
    selectedShifts: shouldClearShiftSelections
      ? new Map()
      : currentState.selectedShifts,
  };
};

export type UseShiftsSelectionActionsOptions = {
  /** Memoized cleanup to run when overlays are closed (e.g. close dropdowns/popovers). */
  closeOverlayCleanup?: () => void;
};

// Hook with helper functions
export const useShiftsSelectionActions = (
  options?: UseShiftsSelectionActionsOptions,
) => {
  const { state, setState, closeOpenOverlays, registerCloseOverlay } =
    useShiftsSelection();

  const { closeOverlayCleanup } = options ?? {};
  useEffect(() => {
    if (closeOverlayCleanup) {
      return registerCloseOverlay(closeOverlayCleanup);
    }
  }, [registerCloseOverlay, closeOverlayCleanup]);

  const toggleAdvancedMode = () => {
    setState(prev => {
      const conciledAdvancedMode =
        prev.isAdvancedMode || prev.isAdvancedModeManual;
      const newState = {
        ...prev,
        isAdvancedMode: !conciledAdvancedMode,
        isAdvancedModeManual: !conciledAdvancedMode,
        // For now when enabling advanced mode, shift selections are cleared
        selectedShifts: !conciledAdvancedMode ? new Map() : prev.selectedShifts,
        excludeShifts: new Map(),
        selectedCells: prev.isAdvancedMode ? [] : prev.selectedCells,
        selectedDates: prev.isAdvancedMode
          ? new Set<string>()
          : prev.selectedDates,
        excludeCells: prev.isAdvancedMode ? [] : prev.excludeCells,
        allRowsSelected: prev.isAdvancedMode ? false : prev.allRowsSelected,
        copiedCells: [],
        pastingCells: new Set<string>(),
      };
      return newState;
    });
  };

  const clearAllSelections = (
    {
      keepDates = false,
      keepCopiedCells = false,
      keepAdvancedMode = false,
      keepSelectedShifts = false,
    }: ClearAllSelectionsParams = {
      keepDates: false,
      keepCopiedCells: false,
      keepAdvancedMode: false,
      keepSelectedShifts: false,
    },
  ) => {
    setState(prev => {
      // No selections after clearing, so calculate advanced mode state
      const advancedModeState = keepAdvancedMode
        ? { isAdvancedMode: prev.isAdvancedMode }
        : calculateAdvancedModeState(false, prev);

      return {
        ...prev,
        selectedDates: keepDates ? prev.selectedDates : new Set<string>(),
        selectedCells: [],
        excludeCells: [],
        excludeShifts: new Map(),
        allRowsSelected: false,
        ...advancedModeState,
        selectedShifts: keepSelectedShifts ? prev.selectedShifts : new Map(), // Overwrite selected shifts from advanced mode state
        copiedCells: keepCopiedCells ? prev.copiedCells : [],
        pastingCells: new Set(),
      };
    });
  };

  const clearCopyPaste = () => {
    setState(prev => ({
      ...prev,
      copiedCells: [],
      pastingCells: new Set(),
    }));
  };

  const clearSelectedCells = () => {
    setState(prev => ({
      ...prev,
      selectedCells: [],
    }));
  };

  const clearShiftSelections = () => {
    setState(prev => ({
      ...prev,
      selectedShifts: new Map(),
      copiedCells: [],
    }));
  };

  const setSelectedDates = (
    dates: Set<string> | ((prev: Set<string>) => Set<string>),
  ) => {
    setState(prev => {
      const newDates =
        typeof dates === 'function' ? dates(prev.selectedDates) : dates;
      const advancedModeState = calculateAdvancedModeState(
        newDates.size > 0,
        prev,
      );

      return {
        ...prev,
        selectedDates: newDates,
        ...advancedModeState,
      };
    });
  };

  const setSelectedCells = (cells: ShiftAssignCellVC[]) => {
    const uniqueCells = cells.filter(
      (cell, index, self) =>
        index ===
        self.findIndex(t => t.userId === cell.userId && t.date === cell.date),
    );
    setState(prev => {
      const advancedModeState = calculateAdvancedModeState(
        uniqueCells.length > 0,
        prev,
      );

      return {
        ...prev,
        selectedCells: uniqueCells,
        ...advancedModeState,
      };
    });
  };

  const setExcludeCells = (
    cells:
      | ShiftAssignCellVC[]
      | ((prev: ShiftAssignCellVC[]) => ShiftAssignCellVC[]),
  ) => {
    setState(prev => {
      const newState = {
        ...prev,
        excludeCells:
          typeof cells === 'function' ? cells(prev.excludeCells) : cells,
      };
      return newState;
    });
  };

  const setViewMode = (mode: WeekMonthPeriod) => {
    setState(prev => ({
      ...prev,
      viewMode: mode,
      selectedCells: [],
      selectedDates: new Set(),
      selectedShifts: new Map(),
      copiedCells: [],
      excludeCells: [],
      excludeShifts: new Map(),
    }));
  };

  const setStartDate = (date: Date) => {
    setState(prev => {
      const newState = {
        ...prev,
        startDate: date,
        selectedDates: new Set<string>(),
        selectedCells: [],
        excludeCells: [],
        excludeShifts: new Map(),
        allRowsSelected: false,
        selectedShifts: new Map(),
        pastingCells: new Set<string>(),
      };
      return newState;
    });
  };

  const setEndDate = (date: Date) => {
    setState(prev => {
      const newState = {
        ...prev,
        endDate: date,
        selectedDates: new Set<string>(),
        selectedCells: [],
        excludeCells: [],
        excludeShifts: new Map(),
        allRowsSelected: false,
        selectedShifts: new Map(),
        pastingCells: new Set<string>(),
      };
      return newState;
    });
  };

  const setAllRowsSelected = (value: boolean) => {
    setState(prev => {
      // Calculate if there are any selections: allRowsSelected OR dates OR cells
      const hasSelections =
        value || prev.selectedDates.size > 0 || prev.selectedCells.length > 0;
      const advancedModeState = calculateAdvancedModeState(hasSelections, prev);

      return {
        ...prev,
        allRowsSelected: value,
        ...advancedModeState,
      };
    });
  };

  /**
   * Toggles the "select all" state. Used by the header checkbox, keyboard
   * shortcut (Cmd/Ctrl+A), and the actions menu. Uses setState directly and
   * applies advanced mode so behaviour is consistent regardless of trigger.
   * @param days - Date keys to select when enabling "select all"
   * @param registers - If provided and empty, no-op (no users in table)
   */
  const toggleSelectAll = (
    days: { key: string }[],
    registers?: ShiftCalendarUser[] | undefined,
  ) => {
    if (registers !== undefined && !registers.length) {
      return;
    }

    setState(prev => {
      if (prev.allRowsSelected && prev.excludeCells.length === 0) {
        const advancedModeState = calculateAdvancedModeState(false, prev);
        return {
          ...prev,
          allRowsSelected: false,
          selectedCells: [],
          selectedDates: new Set<string>(),
          excludeShifts: new Map(),
          ...advancedModeState,
        };
      }

      const nextState = {
        ...prev,
        allRowsSelected: true,
        excludeCells: [],
        selectedCells: [],
        selectedDates: new Set(days.map(d => d.key)),
        selectedShifts: new Map(),
        excludeShifts: new Map(),
      };
      const advancedModeState = calculateAdvancedModeState(true, nextState);
      return { ...nextState, ...advancedModeState };
    });
  };

  const setSelectedShifts = (shifts: Map<number, ShiftSelection>) => {
    setState(prev => {
      const newState = { ...prev, selectedShifts: shifts };
      return newState;
    });
  };

  const toggleShiftSelection = (
    shiftId: number,
    userId: number,
    date: string,
  ) => {
    setState(prev => {
      const newMap = new Map(prev.selectedShifts);
      const isShiftSelected = newMap.has(shiftId);

      if (isShiftSelected) {
        newMap.delete(shiftId);
      } else {
        newMap.set(shiftId, { shiftId, userId, date });
      }

      return { ...prev, selectedShifts: newMap };
    });
  };

  const setExcludeShifts = (
    shifts:
      | Map<number, ShiftSelection>
      | ((prev: Map<number, ShiftSelection>) => Map<number, ShiftSelection>),
  ) => {
    setState(prev => {
      const newState = {
        ...prev,
        excludeShifts:
          typeof shifts === 'function' ? shifts(prev.excludeShifts) : shifts,
      };
      return newState;
    });
  };

  // Only use use to update multiple states at once (all, dates, cells, shifts)
  const updateSelectionState = (updates: Partial<ShiftsSelectionState>) => {
    setState(prev => {
      const newState = { ...prev, ...updates };

      // Calculate hasSelections based on the NEW state (after updates)
      const hasSelections =
        newState.allRowsSelected ||
        newState.selectedDates.size > 0 ||
        newState.selectedCells.length > 0 ||
        newState.selectedShifts.size > 0;

      // Use newState so advancedModeState doesn't overwrite explicit updates
      // (e.g. selectedShifts: new Map()) with stale prev values
      const advancedModeState = calculateAdvancedModeState(
        hasSelections,
        newState,
      );

      return {
        ...newState,
        ...advancedModeState,
      };
    });
  };

  const handleShiftClick = (
    shiftId: number,
    containerCell: ShiftAssignCellVC,
    allShiftsInCell: number[],
    cellSelectionMap: Map<string, boolean>,
  ) => {
    setState(prev => {
      // If shift is excluded, restore it (remove from excludeShifts)
      if (prev.excludeShifts.has(shiftId)) {
        const newExcludeShifts = new Map(prev.excludeShifts);
        newExcludeShifts.delete(shiftId);
        return {
          ...prev,
          excludeShifts: newExcludeShifts,
        };
      }

      const isCellExcluded = prev.excludeCells.some(
        cell =>
          cell.userId === containerCell.userId &&
          cell.date === containerCell.date,
      );

      // If cell is excluded, restore it and add all other shifts to excludeShifts
      if (isCellExcluded) {
        const newExcludeCells = prev.excludeCells.filter(
          cell =>
            cell.userId !== containerCell.userId ||
            cell.date !== containerCell.date,
        );
        const newExcludeShifts = new Map(prev.excludeShifts);
        allShiftsInCell.forEach(id => {
          if (id !== shiftId) {
            newExcludeShifts.set(id, {
              shiftId: id,
              userId: containerCell.userId,
              date: containerCell.date,
            });
          }
        });
        return {
          ...prev,
          excludeCells: newExcludeCells,
          excludeShifts: newExcludeShifts,
        };
      }

      const explicitCellSelection = cellSelectionMap.get(
        `${containerCell.userId}-${containerCell.date}`,
      );

      const isShiftSelected =
        prev.selectedShifts.has(shiftId) ||
        (!prev.excludeShifts.has(shiftId) &&
          (explicitCellSelection ||
            (explicitCellSelection === undefined &&
              prev.selectedDates.has(containerCell.date))));

      const newSelectedShifts = new Map(prev.selectedShifts);

      // If shift is selected (explicitly or implicitly)
      if (isShiftSelected) {
        // Remove from selected if included
        if (newSelectedShifts.has(shiftId)) {
          newSelectedShifts.delete(shiftId);
        }

        // Only add to excludeShifts if there's an implicit selection (cell or date selected)
        const hasImplicitSelection =
          explicitCellSelection ||
          (explicitCellSelection === undefined &&
            prev.selectedDates.has(containerCell.date));

        if (hasImplicitSelection) {
          const newExcludeShifts = new Map(prev.excludeShifts);
          newExcludeShifts.set(shiftId, {
            shiftId,
            userId: containerCell.userId,
            date: containerCell.date,
          });
          return {
            ...prev,
            selectedShifts: newSelectedShifts,
            excludeShifts: newExcludeShifts,
          };
        }

        return {
          ...prev,
          selectedShifts: newSelectedShifts,
        };
      } else {
        newSelectedShifts.set(shiftId, {
          shiftId,
          userId: containerCell.userId,
          date: containerCell.date,
        });
        // Check if all shifts in cell are now selected
        if (allShiftsInCell.every(id => newSelectedShifts.has(id))) {
          // Since all shifts in cell are selected, remove only this cell's shifts and add cell to selectedCells
          const updatedShifts = filterShiftsByIds(
            newSelectedShifts,
            allShiftsInCell,
          );

          return {
            ...prev,
            selectedShifts: updatedShifts,
            selectedCells: [...prev.selectedCells, containerCell],
          };
        }
      }
      // If with this shift all other shifts in cell are selected, add to selectedCells (remove from excludeCells if necessary)
      return { ...prev, selectedShifts: newSelectedShifts };
    });
  };

  const handleSelectCell = (
    cellData: ShiftAssignCellVC,
    allShiftsInCell: number[] = [],
    dates?: string[],
    visibleUserIds?: number[],
  ) => {
    setState(prev => {
      const cellExcluded = prev.excludeCells.some(
        cell => cell.userId === cellData.userId && cell.date === cellData.date,
      );

      // If cell is in excludeCells, restore it
      if (cellExcluded) {
        const newExcludeCells = prev.excludeCells.filter(
          cell =>
            !(cell.userId === cellData.userId && cell.date === cellData.date),
        );

        // Restore any excluded shifts for this cell
        const newExcludeShifts =
          allShiftsInCell.length > 0
            ? filterShiftsByIds(prev.excludeShifts, allShiftsInCell)
            : prev.excludeShifts;
        return {
          ...prev,
          excludeCells: newExcludeCells,
          excludeShifts: newExcludeShifts,
        };
      }

      // Check if this cell has excluded shifts
      const cellHasExcludedShifts =
        allShiftsInCell.length > 0 &&
        allShiftsInCell.some(shiftId => prev.excludeShifts.has(shiftId));

      // Add to selectedCells if:
      // 1. Date is not selected, OR
      // 2. Date is selected but cell has excluded shifts (needs explicit selection)
      let newSelectedCells = prev.selectedCells;
      if (!prev.selectedDates.has(cellData.date) || cellHasExcludedShifts) {
        // Remove if already exists, then add (to avoid duplicates)
        newSelectedCells = [
          ...prev.selectedCells.filter(
            cell =>
              !(cell.userId === cellData.userId && cell.date === cellData.date),
          ),
          cellData,
        ];
      }

      // Filter out redundant shifts within this cell
      const newSelectedShifts = filterShiftsByCells(prev.selectedShifts, [
        cellData,
      ]);

      // Restore any excluded shifts for this cell when selecting
      const newExcludeShifts = filterShiftsByCells(prev.excludeShifts, [
        cellData,
      ]);

      // Skip promotions if cell has excluded shifts (we're explicitly selecting it)
      // Also skip promotions if allRowsSelected is true (we're promoting to all rows)
      const checkForPromotions =
        dates &&
        visibleUserIds &&
        !newExcludeShifts.size &&
        !prev.allRowsSelected;
      const { promotedDates, promotedCells, shouldPromoteToAll } =
        checkForPromotions
          ? checkCellSelectionPromotions(
              cellData,
              newSelectedCells,
              prev.selectedDates,
              dates,
              visibleUserIds,
            )
          : {
              promotedDates: new Set(prev.selectedDates),
              promotedCells: newSelectedCells,
              shouldPromoteToAll: prev.allRowsSelected,
            };

      // Calculate hasSelections based on the NEW state
      const hasSelections =
        shouldPromoteToAll ||
        promotedDates.size > 0 ||
        promotedCells.length > 0 ||
        newSelectedShifts.size > 0;

      const advancedModeState = calculateAdvancedModeState(hasSelections, prev);

      return {
        ...prev,
        selectedDates: promotedDates,
        selectedCells: promotedCells,
        excludeShifts: newExcludeShifts,
        allRowsSelected: shouldPromoteToAll,
        ...advancedModeState,
        // Override selectedShifts after advancedModeState to keep our filtered version
        selectedShifts: newSelectedShifts,
      };
    });
  };

  const handleDeselectCell = (
    cellData: ShiftAssignCellVC,
    allShiftsInCell: number[] = [],
  ) => {
    setState(prev => {
      // If cell is implicitly selected via date or allRows, exclude it
      if (prev.allRowsSelected || prev.selectedDates.has(cellData.date)) {
        const newExcludeCells = [...prev.excludeCells, cellData];

        // Restore any excluded shifts for this cell since the entire cell is now excluded
        const newExcludeShifts =
          allShiftsInCell.length > 0
            ? filterShiftsByIds(prev.excludeShifts, allShiftsInCell)
            : prev.excludeShifts;

        // Also remove from selectedCells if it's explicitly there
        const cellInSelectedCells = prev.selectedCells.some(
          cell =>
            cell.userId === cellData.userId && cell.date === cellData.date,
        );

        const newSelectedCells = cellInSelectedCells
          ? prev.selectedCells.filter(
              cell =>
                !(
                  cell.userId === cellData.userId && cell.date === cellData.date
                ),
            )
          : prev.selectedCells;

        return {
          ...prev,
          excludeCells: newExcludeCells,
          excludeShifts: newExcludeShifts,
          selectedCells: newSelectedCells,
        };
      }

      // Otherwise remove from selectedCells
      const newSelectedCells = prev.selectedCells.filter(
        cell =>
          !(cell.userId === cellData.userId && cell.date === cellData.date),
      );

      // Restore any excluded shifts for this cell when deselecting
      const newExcludeShifts =
        allShiftsInCell.length > 0
          ? filterShiftsByIds(prev.excludeShifts, allShiftsInCell)
          : prev.excludeShifts;

      // Calculate hasSelections based on the NEW state
      const hasSelections =
        prev.allRowsSelected ||
        prev.selectedDates.size > 0 ||
        newSelectedCells.length > 0 ||
        prev.selectedShifts.size > 0;

      const advancedModeState = calculateAdvancedModeState(hasSelections, prev);

      return {
        ...prev,
        selectedCells: newSelectedCells,
        excludeShifts: newExcludeShifts,
        ...advancedModeState,
      };
    });
  };

  const handleDateSelection = (
    date: string,
    isChecked: boolean,
    isIndeterminate: boolean,
    totalDays: number,
    allShiftsInDate: ShiftSelection[] = [],
  ) => {
    setState(prev => {
      if (isChecked) {
        // Deselect the column
        const newDates = new Set(prev.selectedDates);
        newDates.delete(date);

        // Remove cells and shifts under this date from selections
        const newSelectedCells = filterCellsByDate(prev.selectedCells, date);
        const newSelectedShifts = filterShiftsByDate(prev.selectedShifts, date);

        // Remove cells and shifts under this date from exclusions
        const newExcludeCells = filterCellsByDate(prev.excludeCells, date);
        const newExcludeShifts = filterShiftsByDate(prev.excludeShifts, date);

        const newState = {
          ...prev,
          allRowsSelected: false,
          selectedDates: newDates,
          selectedCells: newSelectedCells,
          selectedShifts: newSelectedShifts,
          excludeCells: newExcludeCells,
          excludeShifts: newExcludeShifts,
        };

        const hasSelections =
          prev.allRowsSelected ||
          newDates.size > 0 ||
          newState.selectedCells.length > 0 ||
          newSelectedShifts.size > 0;

        const advancedModeState = calculateAdvancedModeState(
          hasSelections,
          prev,
        );

        return {
          ...newState,
          ...advancedModeState,
        };
      } else if (isIndeterminate) {
        // If there are excluded cells, restore them
        const hasExcludedCells = prev.excludeCells.some(
          cell => cell.date === date,
        );

        if (hasExcludedCells) {
          // Restore excluded cells and clear any excluded shifts for this date
          const newExcludeCells = prev.excludeCells.filter(
            cell => cell.date !== date,
          );

          // Remove all shifts for this date from excludeShifts
          const newExcludeShifts = filterShiftsByIds(
            prev.excludeShifts,
            allShiftsInDate.map(s => s.shiftId),
          );

          return {
            ...prev,
            excludeCells: newExcludeCells,
            excludeShifts: newExcludeShifts,
          };
        } else {
          // No excluded cells, but could be:
          // 1. Date selected with excluded shifts, OR
          // 2. Date not selected but has partial cell selections

          const newExcludeShifts = filterShiftsByIds(
            prev.excludeShifts,
            allShiftsInDate.map(s => s.shiftId),
          );

          const commonNewState = {
            ...prev,
            excludeShifts: newExcludeShifts,
          };

          // If date is not already selected, add it
          if (!prev.selectedDates.has(date)) {
            const newDates = new Set(prev.selectedDates);
            newDates.add(date);

            // Filter out cells for this date since they're now implicitly selected via the date
            const newSelectedCells = filterCellsByDate(
              prev.selectedCells,
              date,
            );
            const newSelectedShifts = filterShiftsByDate(
              prev.selectedShifts,
              date,
            );

            return {
              ...commonNewState,
              selectedDates: newDates,
              selectedCells: newSelectedCells,
              selectedShifts: newSelectedShifts,
            };
          }

          return commonNewState;
        }
      } else {
        // Select the column
        const newDates = new Set(prev.selectedDates);
        newDates.add(date);
        const shouldSelectAllRows = newDates.size === totalDays;

        // Filter out redundant cells and shifts under this date
        const newSelectedCells = filterCellsByDate(prev.selectedCells, date);
        const newSelectedShifts = filterShiftsByDate(prev.selectedShifts, date);

        // Clear exclusions for this date
        const newExcludeCells = filterCellsByDate(prev.excludeCells, date);
        const newExcludeShifts = filterShiftsByDate(prev.excludeShifts, date);

        const hasSelections =
          shouldSelectAllRows ||
          newDates.size > 0 ||
          newSelectedCells.length > 0 ||
          newSelectedShifts.size > 0;

        const advancedModeState = calculateAdvancedModeState(
          hasSelections,
          prev,
        );

        return {
          ...prev,
          selectedDates: newDates,
          selectedCells: newSelectedCells,
          excludeCells: newExcludeCells,
          excludeShifts: newExcludeShifts,
          allRowsSelected: shouldSelectAllRows,
          ...advancedModeState,
          selectedShifts: newSelectedShifts,
        };
      }
    });
  };

  const handleCollaboratorSelection = (
    userId: number,
    isChecked: boolean,
    isIndeterminate: boolean,
    days: { key: string; label: string }[],
    allShiftsInRow: ShiftSelection[] = [],
  ) => {
    setState(prev => {
      if (isChecked) {
        // Deselect the row
        // Remove shifts for this user from selections
        const newSelectedShifts = filterShiftsByCollaborator(
          prev.selectedShifts,
          userId,
        );

        if (prev.allRowsSelected) {
          // Remove from selectedCells and add to excludeCells
          const newExcludeCells = [
            ...filterCellsByCollaborators(prev.excludeCells, [userId]),
            ...days.map(d => ({
              userId,
              date: d.key,
            })),
          ];

          // Clear shift exclusions for cells that are now excluded (redundancy)
          const newExcludeShifts = filterShiftsByCollaborator(
            prev.excludeShifts,
            userId,
          );
          const nextState = {
            ...prev,
            selectedCells: filterCellsByCollaborators(prev.selectedCells, [
              userId,
            ]),
            selectedShifts: newSelectedShifts,
            excludeCells: newExcludeCells,
            excludeShifts: newExcludeShifts,
          };
          const hasSelections =
            nextState.allRowsSelected ||
            nextState.selectedDates.size > 0 ||
            nextState.selectedCells.length > 0 ||
            nextState.selectedShifts.size > 0;
          return {
            ...nextState,
            ...calculateAdvancedModeState(hasSelections, nextState),
          };
        } else if (!prev.allRowsSelected && prev.selectedDates.size > 0) {
          // Only exclude cells for dates that are selected
          const selectedDatesArray = Array.from(prev.selectedDates);
          const newExcludeCells = [
            ...filterCellsByCollaborators(prev.excludeCells, [userId]),
            ...selectedDatesArray.map(d => ({
              userId,
              date: d,
            })),
          ];

          // Clear shift exclusions for cells that are now excluded
          const newExcludeShifts = new Map(prev.excludeShifts);
          for (const [shiftId, shift] of prev.excludeShifts.entries()) {
            if (shift.userId === userId && prev.selectedDates.has(shift.date)) {
              newExcludeShifts.delete(shiftId);
            }
          }

          const nextState = {
            ...prev,
            selectedCells: filterCellsByCollaborators(prev.selectedCells, [
              userId,
            ]),
            selectedShifts: newSelectedShifts,
            excludeCells: newExcludeCells,
            excludeShifts: newExcludeShifts,
          };
          const hasSelections =
            nextState.allRowsSelected ||
            nextState.selectedDates.size > 0 ||
            nextState.selectedCells.length > 0 ||
            nextState.selectedShifts.size > 0;
          return {
            ...nextState,
            ...calculateAdvancedModeState(hasSelections, nextState),
          };
        } else {
          // Just remove from selectedCells
          const nextState = {
            ...prev,
            selectedCells: filterCellsByCollaborators(prev.selectedCells, [
              userId,
            ]),
            selectedShifts: newSelectedShifts,
          };
          const hasSelections =
            nextState.allRowsSelected ||
            nextState.selectedDates.size > 0 ||
            nextState.selectedCells.length > 0 ||
            nextState.selectedShifts.size > 0;
          return {
            ...nextState,
            ...calculateAdvancedModeState(hasSelections, nextState),
          };
        }
      } else if (isIndeterminate) {
        // Toggle from indeterminate to selected
        const newExcludeShifts = filterShiftsByIds(
          prev.excludeShifts,
          allShiftsInRow.map(s => s.shiftId),
        );

        if (prev.allRowsSelected) {
          if (prev.excludeCells.some(cell => cell.userId === userId)) {
            // Restore excluded cells and shifts
            return {
              ...prev,
              excludeCells: prev.excludeCells.filter(
                cell => cell.userId !== userId,
              ),
              excludeShifts: newExcludeShifts,
            };
          } else {
            // Add cells to selection and restore excluded shifts
            return {
              ...prev,
              selectedCells: [
                ...prev.selectedCells,
                ...days.map(d => ({
                  userId,
                  date: d.key,
                })),
              ],
              excludeShifts: newExcludeShifts,
            };
          }
        } else {
          // Not all rows selected - could be:
          // 1. Has excluded cells (need to restore them)
          // 2. Partial cell selection (some cells already in selectedCells)
          // 3. Some dates selected (cells implicitly selected via dates)
          // 4. Has excluded shifts (cells might be fully selected)

          // Check if there are excluded cells for this user
          if (prev.excludeCells.some(cell => cell.userId === userId)) {
            // Restore excluded cells
            return {
              ...prev,
              excludeCells: prev.excludeCells.filter(
                cell => cell.userId !== userId,
              ),
              excludeShifts: newExcludeShifts,
            };
          }

          // No excluded cells, proceed with adding missing cells
          const existingCells = prev.selectedCells.filter(
            cell => cell.userId === userId,
          );

          // Find which days are missing from selectedCells
          const missingDays = days.filter(
            d =>
              !existingCells.some(cell => cell.date === d.key) &&
              !prev.selectedDates.has(d.key),
          );

          // Add only missing cells to avoid duplicates
          const newSelectedCells = [
            ...prev.selectedCells,
            ...missingDays.map(d => ({
              userId,
              date: d.key,
            })),
          ];
          return {
            ...prev,
            selectedCells: newSelectedCells,
            excludeShifts: newExcludeShifts,
          };
        }
      } else {
        // Select the row
        const cellsToAdd = days.map(d => ({
          userId,
          date: d.key,
        }));

        // Filter out redundant shifts within the newly selected cells
        const newSelectedShifts = filterShiftsByCollaborator(
          prev.selectedShifts,
          userId,
        );

        // Clear any excluded cells and shifts for this user
        const newExcludeCells = filterCellsByCollaborators(prev.excludeCells, [
          userId,
        ]);
        const newExcludeShifts = filterShiftsByCollaborator(
          prev.excludeShifts,
          userId,
        );

        if (prev.allRowsSelected) {
          return {
            ...prev,
            excludeCells: newExcludeCells,
            excludeShifts: newExcludeShifts,
            selectedCells: [...prev.selectedCells, ...cellsToAdd],
            selectedShifts: newSelectedShifts,
          };
        } else {
          const newSelectedCells = [...prev.selectedCells, ...cellsToAdd];

          const hasSelections =
            prev.allRowsSelected ||
            prev.selectedDates.size > 0 ||
            newSelectedCells.length > 0 ||
            newSelectedShifts.size > 0;

          const advancedModeState = calculateAdvancedModeState(
            hasSelections,
            prev,
          );

          return {
            ...prev,
            selectedCells: newSelectedCells,
            excludeCells: newExcludeCells,
            excludeShifts: newExcludeShifts,
            ...advancedModeState,
            // Override selectedShifts after advancedModeState to keep our filtered version
            selectedShifts: newSelectedShifts,
          };
        }
      }
    });
  };

  const setCopiedCells = (copiedData: CopiedShiftCell[]) => {
    setState(prev => ({
      ...prev,
      copiedCells: copiedData,
      selectedCells: [],
      selectedShifts: new Map(),
      excludeCells: [],
      excludeShifts: new Map(),
      selectedDates: new Set(),
    }));
  };

  const setPastingCells = (cells: Set<string>) => {
    setState(prev => ({ ...prev, pastingCells: cells }));
  };

  return {
    ...state,

    // Actions
    closeOpenOverlays,
    registerCloseOverlay,
    toggleAdvancedMode,
    clearAllSelections,
    clearShiftSelections,
    setSelectedDates,
    setSelectedCells,
    setExcludeCells,
    setViewMode,
    setStartDate,
    setEndDate,
    setAllRowsSelected,
    toggleSelectAll,
    setSelectedShifts,
    toggleShiftSelection,
    setExcludeShifts,
    updateSelectionState,
    handleShiftClick,
    handleSelectCell,
    handleDeselectCell,
    handleDateSelection,
    handleCollaboratorSelection,
    setCopiedCells,
    setPastingCells,
    clearCopyPaste,
    clearSelectedCells,
  };
};
