import {
  type MouseEvent,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import { useFormContext } from 'react-hook-form';
import {
  type FetchNextPageOptions,
  type InfiniteQueryObserverResult,
} from 'react-query';

import { useModal } from '@material-hu/hooks/useModal';
import useVirtualizer from '@material-hu/hooks/useVirtualizer';
import {
  IconCopy,
  IconExclamationCircle,
  IconTrash,
} from '@material-hu/icons/tabler';
import Divider from '@material-hu/mui/Divider';
import Popover from '@material-hu/mui/Popover';
import Stack from '@material-hu/mui/Stack';
import Typography from '@material-hu/mui/Typography';

import Avatar from '@material-hu/components/design-system/Avatar';
import Checkbox from '@material-hu/components/design-system/Checkbox/Checkbox';
import Dialog from '@material-hu/components/design-system/Dialog';
import MenuItem from '@material-hu/components/design-system/Menu/components/MenuItem';
import Spinner from '@material-hu/components/design-system/ProgressIndicators/Spinner';
import Table from '@material-hu/components/design-system/Table';
import TableCell from '@material-hu/components/design-system/Table/components/TableCell';
import TableContainer from '@material-hu/components/design-system/Table/components/TableContainer';
import TableRow from '@material-hu/components/design-system/Table/components/TableRow';
import Title from '@material-hu/components/design-system/Title';

import { queryClient } from 'src/config/react-query';
import { logEvent } from 'src/config/amplitude';
import { useAuth } from 'src/contexts/JWTContext';
import { getUserListByIds } from 'src/services/users';
import {
  type Shift,
  type SelectedAssignedShift,
  type ShiftAssignCellVC,
  type ShiftCalendarDay,
  type ShiftCalendarUser,
  type ShiftTimeSlot,
  type TimeOffShiftPopoverInfo,
} from 'src/types/shifts';
import { isAdmin } from 'src/utils/permissions';
import {
  calculateTimeSpanInHours,
  parseHoursAndMinutesDisplay,
} from 'src/utils/timeTracking';
import {
  EventName,
  ShiftManagementMenuOption,
  ShiftManagementMenuSource,
} from 'src/types/amplitude';
import { useLokaliseTranslation } from 'src/utils/i18n';

type ShiftContextMenuInfo = {
  name: string | null;
  timeSlots: ShiftTimeSlot[];
  isDayOff?: boolean;
};

import EmptyState from 'src/components/HuEmptyState';
import InfiniteTable from 'src/components/InfiniteTable';

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

import { TABLE_HEADERS_IDS } from '../../constants';
import { useShiftsSelectionActions } from '../../contexts/ShiftsSelectionContext';
import { shiftFormFields } from '../../form';
import { useShiftBulkDelete } from '../../hooks/useShiftBulkDelete';
import { shiftsKeys } from '../../queries';
import TimeOffPopover from '../TimeOffPopover';

import CollaboratorCell from './CollaboratorCell';
import {
  BASE_ROW_HEIGHT,
  FIXED_ROW_HEIGHT_ESTIMATE,
  SELECTION_COL_WIDTH,
  SKELETON_ROWS,
  STICKY_COLUMN_SX,
} from './constants';
import DateCell from './DateCell';
import ShiftsTableSkeleton from './ShiftsTableSkeleton';
import { getRowCheckboxState, getShiftsForCollaborator } from './utils';
import { type OptimisticMove } from '../../hooks/useShiftDragDrop';
import { type ShiftDragDropHandlers } from '../../types';

type Props = {
  isLoading: boolean;
  isFetchingNextPage: boolean;
  hasNextPage: boolean;
  fetchNextPage: (
    options?: FetchNextPageOptions,
  ) => Promise<InfiniteQueryObserverResult<unknown, unknown>>;
  totalRegisters: number;
  registers: ShiftCalendarUser[];
  showAssignmentDrawer: (drawerProps: Partial<unknown>) => void;
  onCopy: (standaloneShift?: {
    shiftId: number;
    userId: number;
    date: string;
  }) => void;
  onPaste: (targetCell?: ShiftAssignCellVC) => void;
  userCount: string;
  horizontalScrollRef: (element: HTMLElement | null) => void;
  days: { key: string; label: string }[];
  selectionActive: boolean;
  appliedFilters: boolean;
  dragOverCell: ShiftAssignCellVC | null;
  draggingShiftIds: Set<number>;
  optimisticMove: OptimisticMove | null;
  /** Maps `${userId}-${date}` → number of shifts that will land there during an active drag. */
  dragPreviewMap: Map<string, number>;
} & ShiftDragDropHandlers;

/**
 * Applies the optimistic move overlay on top of the real calendar data.
 * Used during the 5-second pending-move window to show the shift as
 * "already moved" without modifying the React-Query cache.
 */
const applyOptimisticOverride = (
  original: ShiftCalendarDay | undefined,
  optimisticMove: OptimisticMove | null,
  userId: number,
  date: string,
): ShiftCalendarDay | undefined => {
  if (!optimisticMove || !original) return original;

  const key = `${userId}-${date}`;
  const removalIds = optimisticMove.removals.get(key);
  const additionShifts = optimisticMove.additions.get(key);

  if (!removalIds && !additionShifts) return original;

  const filteredShifts: Shift[] = removalIds
    ? (original.shifts ?? []).filter(
        s => !removalIds.has(s.shiftAssignedShiftId),
      )
    : (original.shifts ?? []);

  // Deduplicate: skip any overlay addition whose shift is already present in
  // filteredShifts. This handles the race where a background refetch brings
  // confirmed server data (shift already in target cell) while the committed
  // overlay is still active — without deduplication we'd get two cards.
  const filteredIds = new Set(filteredShifts.map(s => s.shiftAssignedShiftId));
  const finalShifts: Shift[] = additionShifts
    ? [
        ...filteredShifts,
        ...additionShifts.filter(s => !filteredIds.has(s.shiftAssignedShiftId)),
      ]
    : filteredShifts;

  return { ...original, shifts: finalShifts };
};

const ShiftContextMenuHeader = ({ info }: { info: ShiftContextMenuInfo }) => {
  const { t } = useLokaliseTranslation(['shifts', 'time_tracker']);
  const totalHours = info.timeSlots.reduce(
    (acc, slot) => acc + calculateTimeSpanInHours(slot.startTime, slot.endTime),
    0,
  );
  return (
    <>
      <Stack
        sx={{
          flexDirection: 'row',
          justifyContent: 'space-between',
          alignItems: 'center',
          px: 2,
          py: 1.5,
          gap: 2,
        }}
      >
        <Stack>
          <Typography
            variant="globalS"
            sx={{ fontWeight: 'fontWeightSemiBold' }}
          >
            {info.isDayOff
              ? t('time_tracker:day_off')
              : info.name || t('time_tracker:unnamed_shift')}
          </Typography>
          {info.timeSlots.length > 0 && (
            <Typography
              variant="globalXS"
              sx={{ color: ({ palette }) => palette.new.text.neutral.lighter }}
            >
              {t(...parseHoursAndMinutesDisplay(totalHours, true))}
            </Typography>
          )}
        </Stack>
        {info.timeSlots.length > 0 && (
          <Stack sx={{ alignItems: 'flex-end' }}>
            {info.timeSlots.map(slot => (
              <Typography
                key={`${slot.startTime}-${slot.endTime}`}
                variant="globalS"
                sx={{ whiteSpace: 'nowrap' }}
              >
                {slot.startTime} - {slot.endTime}
              </Typography>
            ))}
          </Stack>
        )}
      </Stack>
      <Divider />
    </>
  );
};

const ShiftsTable = ({
  isLoading,
  isFetchingNextPage,
  hasNextPage,
  fetchNextPage,
  totalRegisters,
  registers,
  showAssignmentDrawer,
  onCopy,
  onPaste,
  horizontalScrollRef,
  days,
  selectionActive,
  appliedFilters,
  onDragStart,
  onDragOver,
  onDragLeave,
  onDrop,
  onDragEnd,
  dragOverCell,
  draggingShiftIds,
  optimisticMove,
  dragPreviewMap,
}: Props) => {
  const [shiftActionsAnchorEl, setShiftActionsAnchorEl] =
    useState<HTMLElement | null>(null);
  const [selectedAssignedShift, setSelectedAssignedShift] =
    useState<SelectedAssignedShift | null>(null);
  const [selectedShiftContextInfo, setSelectedShiftContextInfo] =
    useState<ShiftContextMenuInfo | null>(null);
  const [timeOffPopoverAnchor, setTimeOffPopoverAnchor] =
    useState<Element | null>(null);
  const [timeOffPopoverInfo, setTimeOffPopoverInfo] =
    useState<TimeOffShiftPopoverInfo | null>(null);
  const { t } = useLokaliseTranslation([
    'shifts',
    'time_tracking_common',
    'time_tracker',
  ]);

  const { permissions } = useAuth();
  const form = useFormContext();
  const scrollElementRef = useRef<HTMLElement | null>(null);

  const closeOverlayCleanup = useCallback(() => {
    setShiftActionsAnchorEl(null);
    setSelectedAssignedShift(null);
    setSelectedShiftContextInfo(null);
    setTimeOffPopoverAnchor(null);
    setTimeOffPopoverInfo(null);
  }, []);
  const {
    isAdvancedMode,
    selectedDates,
    selectedCells,
    excludeCells,
    setSelectedCells,
    selectedShifts,
    setSelectedShifts,
    toggleShiftSelection,
    excludeShifts,
    handleShiftClick,
    handleSelectCell,
    handleDeselectCell,
    handleCollaboratorSelection,
    copiedCells,
    closeOpenOverlays,
    viewMode,
  } = useShiftsSelectionActions({ closeOverlayCleanup });

  const isMonthlyView = viewMode === WeekMonthPeriod.MONTHLY;

  // Switching views unmounts the anchored shift card, leaving the popover open
  // with a stale anchor (MUI then renders it at the top-left corner). Close it.
  useEffect(() => {
    closeOverlayCleanup();
  }, [viewMode, closeOverlayCleanup]);

  // Calculate once per table render instead of per cell
  const copiedShiftIds = useMemo(() => {
    if (!copiedCells.length) return null;
    return new Set(copiedCells.flatMap(cell => cell.shiftIds));
  }, [copiedCells]);

  const handleCollabActionsClick = (userId: string | number) => {
    // Add the collaborator to selection if not already selected (for visual feedback)
    if (!selectedCells.some(cell => cell.userId === Number(userId))) {
      setSelectedCells([
        ...selectedCells,
        ...days.map(d => ({ userId: Number(userId), date: d.key })),
      ]);
    }
  };

  const handleShiftActionsClick = (
    event: MouseEvent<HTMLButtonElement>,
    assignedShiftData: SelectedAssignedShift,
  ) => {
    event.preventDefault();
    event.stopPropagation();
    setSelectedAssignedShift(assignedShiftData);
    setShiftActionsAnchorEl(event.currentTarget);

    // Add the shift to selection if not already selected (for visual feedback)
    if (!selectedShifts.has(assignedShiftData.shiftId)) {
      toggleShiftSelection(
        assignedShiftData.shiftId,
        assignedShiftData.userId,
        assignedShiftData.date,
      );
    }
  };

  const handleShiftContextMenu = (
    event: MouseEvent<HTMLElement>,
    assignedShiftData: SelectedAssignedShift,
    shiftInfo: ShiftContextMenuInfo,
  ) => {
    event.preventDefault();
    event.stopPropagation();
    logEvent(EventName.SHIFT_MANAGEMENT_RIGHT_CLICK_SHIFT, {
      user_admin: isAdmin(permissions),
    });
    setSelectedAssignedShift(assignedShiftData);
    setSelectedShiftContextInfo(shiftInfo);
    setShiftActionsAnchorEl(event.currentTarget);

    if (!selectedShifts.has(assignedShiftData.shiftId)) {
      toggleShiftSelection(
        assignedShiftData.shiftId,
        assignedShiftData.userId,
        assignedShiftData.date,
      );
    }
  };

  const { executeBulkDelete, executeIndividualDelete, isDeletingShifts } =
    useShiftBulkDelete({
      onSuccess: () => {
        setSelectedAssignedShift(null);
        closeRemoveAssignmentModal();
      },
    });

  const handleDeleteShifts = () => {
    const shiftIds = Array.from(selectedShifts.keys());

    // If multiple shifts are selected, use bulk delete
    if (shiftIds.length > 1) {
      executeBulkDelete();
    } else {
      // Single shift deletion
      executeIndividualDelete({
        shiftAssignedDayId: selectedAssignedShift?.dayId || 0,
        shiftAssignedShiftId: selectedAssignedShift?.shiftId || 0,
      });
    }
  };

  const {
    modal: removeAssignmentModal,
    showModal: showRemoveAssignmentModal,
    closeModal: closeRemoveAssignmentModal,
  } = useModal(
    Dialog,
    { fullWidth: true, maxWidth: 'sm' },
    {
      title:
        selectedShifts.size > 1
          ? t('shifts:remove_modal.bulk_title')
          : t('shifts:remove_modal.single_title'),
      textBody:
        selectedShifts.size > 1
          ? t('shifts:remove_modal.bulk_description')
          : t('shifts:remove_modal.single_description'),
      primaryButtonProps: {
        children: t('remove_shift', { count: selectedShifts.size }),
        loading: isDeletingShifts,
        onClick: handleDeleteShifts,
      },
      secondaryButtonProps: {
        children: t('general:cancel'),
        onClick: () => {
          closeRemoveAssignmentModal();
        },
      },
    },
  );

  const CELL_ACTION_OPTIONS = [
    {
      text: t('remove_from_cell'),
      icon: IconTrash,
      description: t('remove_from_cell_description', {
        count: selectedShifts.size,
      }),
      onClick: () => {
        logEvent(EventName.SHIFT_MANAGEMENT_MENU_USED, {
          option_used: ShiftManagementMenuOption.REMOVE_SHIFT,
          menu: ShiftManagementMenuSource.SHIFT,
        });
        showRemoveAssignmentModal();
        setShiftActionsAnchorEl(null);
      },
    },
    {
      text: t('copy_shift', { count: selectedShifts.size }),
      icon: IconCopy,
      onClick: () => {
        logEvent(EventName.SHIFT_MANAGEMENT_MENU_USED, {
          option_used: ShiftManagementMenuOption.COPY,
          menu: ShiftManagementMenuSource.SHIFT,
        });
        onCopy();
        setShiftActionsAnchorEl(null);
      },
    },
  ];

  // Map of cell selection status for each cell (if null, cell is not selected)
  const cellSelectionMap = useMemo(() => {
    const map = new Map<string, boolean>();

    selectedCells.forEach(cell => {
      map.set(`${cell.userId}-${cell.date}`, true);
    });
    excludeCells.forEach(cell => {
      map.set(`${cell.userId}-${cell.date}`, false);
    });

    return map;
  }, [selectedCells, excludeCells]);

  const handleShiftClickWrapper = useCallback(
    (
      shiftId: number,
      containerCellData: ShiftAssignCellVC,
      allShiftsInCell: number[],
    ) => {
      handleShiftClick(
        shiftId,
        containerCellData,
        allShiftsInCell,
        cellSelectionMap,
      );
    },
    [handleShiftClick, cellSelectionMap],
  );

  const handleSelectCellWrapper = useCallback(
    (cellData: ShiftAssignCellVC, allShiftsInCell: number[]) => {
      const visibleUserIds = registers?.map(r => r.userId) || [];
      const dateKeys = days.map(d => d.key);
      handleSelectCell(cellData, allShiftsInCell, dateKeys, visibleUserIds);
    },
    [handleSelectCell, registers, days],
  );

  const handleSelectAllShiftsInCell = useCallback(
    (dateInfo: ShiftCalendarDay, userId: number, date: string) => {
      // Select all assignments in cell (exclude TO requests)
      const shiftsInCell = dateInfo.shifts.map(
        (s: { shiftAssignedShiftId: number }) => s.shiftAssignedShiftId,
      );

      const allShiftsSelected = shiftsInCell.every((id: number) =>
        selectedShifts.has(id),
      );

      if (allShiftsSelected) {
        // Deselect all shifts from this cell
        const newMap = new Map(selectedShifts);
        shiftsInCell.forEach((id: number) => {
          newMap.delete(id);
        });
        setSelectedShifts(newMap);
      } else {
        // Select all shifts from this cell
        const newMap = new Map(selectedShifts);
        shiftsInCell.forEach((id: number) => {
          newMap.set(id, { shiftId: id, userId, date });
        });
        setSelectedShifts(newMap);
      }
    },
    [selectedShifts, setSelectedShifts],
  );

  const handleCellClick = (collaborator: ShiftCalendarUser, date: string) => {
    if (!isAdvancedMode) {
      queryClient.prefetchQuery(
        shiftsKeys.collaborators(collaborator.userId),
        () => getUserListByIds([collaborator.userId]),
      );
      const dateToSet = new Date(`${date}T00:00:00`);
      form.setValue(shiftFormFields.date, dateToSet);
      form.setValue(shiftFormFields.collaborators, [
        {
          value: collaborator.userId,
          label: collaborator.fullName,
          userInfo: {
            id: collaborator.userId,
            fullName: collaborator.fullName,
            profilePicture: collaborator.profilePicture,
            employeeInternalId: collaborator.employeeInternalId,
          },
        },
      ]);
      showAssignmentDrawer({});
    }
  };

  const handleTimeOffClick = (
    event: MouseEvent<Element>,
    info: TimeOffShiftPopoverInfo,
  ) => {
    setTimeOffPopoverAnchor(event.currentTarget);
    setTimeOffPopoverInfo(info);
  };

  // Estimates the row height considering the number of shifts per row
  const estimateRowSize = useCallback(
    (index: number) => {
      const collaborator = registers[index];
      if (!collaborator) return FIXED_ROW_HEIGHT_ESTIMATE;

      // Get max shifts in any day for this collaborator
      const maxShifts = Object.values(collaborator.days).reduce(
        (max, day) => Math.max(max, day.shifts?.length || 0),
        1,
      );

      // Each shift adds approximately 51px, with base height of 93px
      return (
        BASE_ROW_HEIGHT + (maxShifts - 1) * 51 + Math.floor(maxShifts / 2) * 8
      );
    },
    [registers],
  );

  const { virtualRows, rowVirtualizer } = useVirtualizer({
    scrollElementRef,
    registers,
    virtualizerOptions: {
      count: registers.length,
      estimateSize: estimateRowSize,
      getScrollElement: () => scrollElementRef.current,
      measureElement: element => element.getBoundingClientRect().height,
      enabled: !!registers?.length,
    },
    withWindowScroll: true,
    hasNextPage,
    isFetchingNextPage,
    fetchNextPage,
  });

  const selectionCols = selectionActive ? 1 : 0;
  const columnCount = selectionCols + 1 + days.length;

  return (
    <>
      {removeAssignmentModal}
      <Stack>
        <TableContainer
          aria-describedby={TABLE_HEADERS_IDS.tableHeader}
          containerRef={horizontalScrollRef}
          sx={{
            borderTopLeftRadius: 0,
            borderTopRightRadius: 0,
            border: 0,
          }}
        >
          <Table
            id={TABLE_HEADERS_IDS.tableBody}
            aria-label={t('shifts:shifts_table_body')}
            sx={{
              tableLayout: 'fixed',
              minWidth: 1200,
              // Optimize rendering performance for virtualized rows
              '& tbody tr': {
                willChange: 'auto',
                contain: 'layout style paint',
              },
            }}
            aria-colcount={columnCount}
            aria-rowcount={totalRegisters}
            stickyHeader // Although this table does not have a header the prop is necessary to maintain borders on horizontal scroll
          >
            {/* Colgroup is used to fix the width of the collaborator column since scrolling past virtualized rows causes the columns to shift */}
            <colgroup>
              {selectionActive && (
                <col style={{ width: `${SELECTION_COL_WIDTH}px` }} />
              )}
              <col style={{ width: STICKY_COLUMN_SX.width }} />
              {days.map(day => (
                <col
                  key={day.key}
                  style={isMonthlyView ? { width: '67px' } : undefined}
                />
              ))}
            </colgroup>
            <InfiniteTable
              fetchNextPage={fetchNextPage}
              // Will trigger fetch from virtualizer
              hasNextPage={false}
              // Will apply the spinner outside the table so it does not affect the virtualize measurements
              isFetchingNextPage={false}
              isEmpty={!totalRegisters}
              isLoading={isLoading}
              totalColumns={columnCount}
            >
              {isLoading && (
                <ShiftsTableSkeleton
                  selectionActive={selectionActive}
                  dayCount={days.length}
                />
              )}
              {!isLoading && !registers?.length && (
                <TableRow>
                  <TableCell colSpan={columnCount}>
                    <EmptyState
                      titleProps={{
                        title: t('no_collaborators_title'),
                        description: appliedFilters
                          ? t('no_collaborators_description_filtered')
                          : t('no_collaborators_description'),
                        variant: 'M',
                      }}
                      avatarProps={{
                        Icon: IconExclamationCircle,
                        color: 'default',
                      }}
                      sx={{
                        height: FIXED_ROW_HEIGHT_ESTIMATE * SKELETON_ROWS - 32,
                      }}
                    />
                  </TableCell>
                </TableRow>
              )}
              {!isLoading && virtualRows.length > 0 && (
                <>
                  {/* Spacer row before visible items */}
                  {virtualRows[0] && virtualRows[0].start > 0 && (
                    <TableRow sx={{ height: `${virtualRows[0].start}px` }}>
                      <TableCell colSpan={columnCount} />
                    </TableRow>
                  )}
                  {/* Visible rows */}
                  {virtualRows.map(virtualRow => {
                    const collaborator = registers?.[virtualRow.index];
                    if (!collaborator) return null;

                    const allShiftIdsInRow = getShiftsForCollaborator(
                      collaborator.userId,
                      days,
                      registers,
                    );

                    const {
                      checked: rowChecked,
                      indeterminate: rowIndeterminate,
                    } = getRowCheckboxState({
                      collaborator,
                      days,
                      selectedDates,
                      selectedCells,
                      excludeCells,
                      excludeShifts,
                      rowShifts: allShiftIdsInRow,
                    });

                    return (
                      <TableRow
                        key={collaborator.userId}
                        data-index={virtualRow.index}
                        ref={rowVirtualizer.measureElement}
                        sx={{
                          // Cell padding (py: 1.5 = 24px) + IconButton (40px) + 1px border = 65px
                          height: 65,
                          '&:hover': {
                            backgroundColor: ({ palette }) =>
                              palette.new.background.elements.default,
                          },
                        }}
                      >
                        {selectionActive && (
                          <TableCell
                            align="center"
                            headers={TABLE_HEADERS_IDS.selectionHeader}
                            sx={{
                              backgroundColor: ({ palette }) =>
                                rowChecked
                                  ? palette.new.background.layout.brand
                                  : palette.new.background.layout.tertiary,
                              left: 0,
                              position: 'sticky',
                              width: SELECTION_COL_WIDTH,
                              zIndex: 1,
                            }}
                          >
                            <Checkbox
                              checked={rowChecked}
                              indeterminate={rowIndeterminate}
                              onClick={e => {
                                e.stopPropagation();
                                closeOpenOverlays();
                                handleCollaboratorSelection(
                                  collaborator.userId,
                                  rowChecked,
                                  rowIndeterminate,
                                  days,
                                  allShiftIdsInRow,
                                );
                              }}
                            />
                          </TableCell>
                        )}
                        <CollaboratorCell
                          register={collaborator}
                          onCollabActionsClick={handleCollabActionsClick}
                          selected={rowChecked}
                          isRowSelectionActive={selectionActive}
                          showRadioButton={!selectionActive}
                          collapseMarginLeft={isAdvancedMode}
                        />
                        {days.map(d => {
                          // Apply optimistic move overlay (during 5s pending-move window)
                          const dateInfo = applyOptimisticOverride(
                            collaborator.days[d.key],
                            optimisticMove,
                            collaborator.userId,
                            d.key,
                          );
                          const cellKey = `${collaborator.userId}-${d.key}`;
                          const allShiftsInCellSelected =
                            isAdvancedMode && dateInfo?.shifts
                              ? dateInfo.shifts.every(
                                  shift =>
                                    selectedShifts.has(
                                      shift.shiftAssignedShiftId,
                                    ) &&
                                    !excludeShifts.has(
                                      shift.shiftAssignedShiftId,
                                    ),
                                ) && dateInfo.shifts.length > 0
                              : false;

                          const hasExcludedShifts =
                            dateInfo?.shifts?.some(shift =>
                              excludeShifts.has(shift.shiftAssignedShiftId),
                            ) || false;

                          const explicitCellSelection =
                            cellSelectionMap.get(cellKey);

                          // Cell is selected if:
                          // 1. Explicitly selected via cellSelectionMap, OR
                          // 2. No explicit selection AND (date selected OR all shifts selected)
                          // Cell is NOT selected if it has any excluded shifts
                          const isCellSelected =
                            allShiftsInCellSelected ||
                            (!hasExcludedShifts &&
                              (explicitCellSelection ||
                                (explicitCellSelection === undefined &&
                                  selectedDates.has(d.key))));

                          const isDragOver =
                            dragOverCell?.userId === collaborator.userId &&
                            dragOverCell?.date === d.key;

                          // Number of shift indicator boxes to show in this cell
                          // (> 0 for any cell where dragged shifts would land)
                          const dragPreviewShiftCount =
                            dragPreviewMap.get(cellKey) ?? 0;

                          return (
                            <DateCell
                              key={cellKey}
                              dateInfo={dateInfo}
                              showAssignmentDrawer={() => {
                                handleCellClick(collaborator, d.key);
                              }}
                              cellData={{
                                userId: collaborator.userId,
                                date: d.key,
                              }}
                              onAdvancedShiftClick={handleShiftClickWrapper}
                              onShiftClick={toggleShiftSelection}
                              onSelectCell={handleSelectCellWrapper}
                              onDeselectCell={handleDeselectCell}
                              onSelectAllShiftsInCell={(
                                shiftDateInfo: ShiftCalendarDay,
                              ) =>
                                handleSelectAllShiftsInCell(
                                  shiftDateInfo,
                                  collaborator.userId,
                                  d.key,
                                )
                              }
                              onMenuOpen={handleShiftActionsClick}
                              onShiftContextMenu={handleShiftContextMenu}
                              isCellSelected={isCellSelected}
                              cellSelectionMap={cellSelectionMap}
                              onTimeOffClick={handleTimeOffClick}
                              copiedShiftIds={copiedShiftIds}
                              onPaste={onPaste}
                              // Drag & drop
                              onDragStart={onDragStart}
                              onDragOver={onDragOver}
                              onDragLeave={onDragLeave}
                              onDrop={onDrop}
                              onDragEnd={onDragEnd}
                              isDragOver={isDragOver}
                              draggingShiftIds={draggingShiftIds}
                              dragPreviewShiftCount={dragPreviewShiftCount}
                            />
                          );
                        })}
                      </TableRow>
                    );
                  })}
                  {/* Spacer row after visible items */}
                  {virtualRows.length > 0 && (
                    <TableRow
                      style={{
                        height: `${
                          rowVirtualizer.getTotalSize() -
                          (virtualRows[virtualRows.length - 1]?.end || 0)
                        }px`,
                      }}
                    >
                      <TableCell
                        colSpan={columnCount}
                        sx={{ p: 0, border: 0 }}
                      />
                    </TableRow>
                  )}
                </>
              )}
            </InfiniteTable>
          </Table>
          {isFetchingNextPage && (
            <Stack sx={{ py: 2 }}>
              <Spinner centered />
            </Stack>
          )}
        </TableContainer>
        {/* Shift/Day-off actions menu */}
        <Popover
          anchorEl={shiftActionsAnchorEl}
          open={!!shiftActionsAnchorEl}
          onClose={() => {
            setShiftActionsAnchorEl(null);
            setSelectedAssignedShift(null);
            setSelectedShiftContextInfo(null);
          }}
          sx={{ minWidth: 320, zIndex: 1 }}
        >
          {isMonthlyView && selectedShiftContextInfo && (
            <ShiftContextMenuHeader info={selectedShiftContextInfo} />
          )}
          {CELL_ACTION_OPTIONS.map(({ icon, onClick, text, description }) => (
            <MenuItem
              key={text}
              onClick={onClick}
            >
              <Avatar
                Icon={icon}
                sx={{
                  backgroundColor: ({ palette }) =>
                    palette.new.background.elements.grey,
                  color: ({ palette }) =>
                    palette.new.action.button.text.tertiary.default,
                  mr: 1,
                }}
              ></Avatar>
              <Title
                title={text}
                description={description || undefined}
                variant="S"
                sx={{ whiteSpace: 'normal' }}
              />
            </MenuItem>
          ))}
        </Popover>
      </Stack>
      <TimeOffPopover
        anchor={timeOffPopoverAnchor}
        closePopover={() => {
          setTimeOffPopoverAnchor(null);
        }}
        timeOffInfo={timeOffPopoverInfo}
      />
    </>
  );
};

export default ShiftsTable;
