import type React from 'react';
import { createContext, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import type { FocusData, TableContextType, TDataWithGroupId, UseFocusHandlerProps } from '../types';
import type { Table } from '@tanstack/react-table';
import { useTableContext } from '../table-context/table-context-provider';
import { CellTypes } from '../constants';

const FocusDataContext = createContext<FocusData>({
  columnIndex: 0,
  rowId: null,
  cellType: CellTypes.headerCell,
  groupId: null,
  rowIndex: 0,
});
const FocusDispatchContext = createContext<(data: FocusData) => void>(() => {});

interface FocusHandlerProviderProps<TData> {
  table: Table<TData>;
  children: React.ReactNode;
  groupData: TableContextType<TData>['groupData'];
}
export default function FocusHandlerContextProvider<TData>(
  props: FocusHandlerProviderProps<TData>,
) {
  const [focusedCell, setFocusedCell] = useState<FocusData>({
    columnIndex: 0,
    rowId: null,
    cellType: CellTypes.headerCell,
    groupId: null,
    rowIndex: 0,
  });

  const { tableRef, allowBulkGroupSelection, hasFooter } = useTableContext();

  // Inject the ref into the table element

  function updateData(data: FocusData) {
    setFocusedCell(data);
  }

  // Table Keyboard Navigation Logic
  useEffect(() => {
    const handleKeyDown = (e: KeyboardEvent) => {
      if (!tableRef.current || !tableRef.current.contains(e.target as Node)) {
        return;
      }

      const activeElement = document.activeElement;
      const isInsideCustomSelect = (el: Element | null) => {
        return el && (el.closest('[role="listbox"]') || el.closest('[role="combobox"]'));
      };

      if (
        activeElement?.tagName === 'SELECT' ||
        activeElement?.tagName === 'INPUT' ||
        isInsideCustomSelect(activeElement)
      ) {
        return;
      }

      const { rowId, columnIndex, cellType, groupId, rowIndex } = focusedCell;

      const rows = props.table.getRowModel().rows;
      const tableHeaders = props.table.getHeaderGroups()[0]?.headers || [];

      if (['ArrowUp', 'ArrowDown', 'ArrowLeft', 'ArrowRight'].includes(e.key)) {
        e.preventDefault();

        let newrowId = rowId;
        let newColumnIndex = columnIndex;
        let newCellType = cellType;
        let newGroupId = groupId;
        let newRowIndex = rowIndex;

        const isGroupView = props.groupData.length > 0;

        switch (e.key) {
          case 'ArrowUp':
            if (cellType === CellTypes.headerCell) {
              return;
            }
            if (hasFooter && cellType === CellTypes.footerCell) {
              newCellType = CellTypes.dataCell;
              newGroupId = null;
              newRowIndex = rows.length - 1;
              newrowId = rows[rows.length - 1].id;
            } else if (cellType === CellTypes.groupHeaderCell) {
              const index = rows.findIndex(
                ({ original }) =>
                  (original as TDataWithGroupId<TData>).editableTableGroupId === groupId,
              );

              if (index === 0) {
                newCellType = CellTypes.headerCell;
                newrowId = null;
                newGroupId = null;
                newRowIndex = -1;
              } else {
                newCellType = CellTypes.dataCell;
                newrowId = rows[index - 1].id;
                newGroupId = null;
                newRowIndex = index - 1;
              }
            } else if (cellType === CellTypes.dataCell) {
              const index = rows.findIndex((row) => row.id === rowId);

              if (isGroupView) {
                if (
                  index === 0 ||
                  (rows[index].original as TDataWithGroupId<TData>).editableTableGroupId !==
                    (rows[index - 1]?.original as TDataWithGroupId<TData>)?.editableTableGroupId
                ) {
                  newGroupId = (rows[index].original as TDataWithGroupId<TData>)
                    .editableTableGroupId;
                  newrowId = null;
                  newCellType = CellTypes.groupHeaderCell;
                  newRowIndex = -1;
                } else {
                  newrowId = rows[Math.min(rows.length - 1, index - 1)].id;
                  newRowIndex = Math.min(rows.length - 1, index - 1);
                }
              } else {
                if (index === 0) {
                  newCellType = CellTypes.headerCell;
                  newGroupId = null;
                  newrowId = null;
                  newRowIndex = -1;
                } else {
                  newCellType = CellTypes.dataCell;
                  newGroupId = null;
                  newrowId = rows[Math.max(0, index - 1)].id;
                  newRowIndex = index - 1;
                }
              }
            }
            break;
          case 'ArrowDown':
            if (cellType === CellTypes.headerCell) {
              if (isGroupView) {
                newCellType = CellTypes.groupHeaderCell;
                newGroupId = (rows[0].original as TDataWithGroupId<TData>).editableTableGroupId;
                newrowId = null;
                newRowIndex = -1;
              } else {
                newGroupId = null;
                newCellType = CellTypes.dataCell;
                newrowId = rows[Math.min(rows.length - 1, 0)].id;
                newRowIndex = Math.min(rows.length - 1, 0);
              }
            } else if (cellType === CellTypes.groupHeaderCell) {
              newCellType = CellTypes.dataCell;
              const index = rows.findIndex(
                ({ original }) =>
                  (original as TDataWithGroupId<TData>).editableTableGroupId === groupId,
              );
              newrowId = rows[index].id;
              newGroupId = null;
              newRowIndex = index;
            } else if (cellType === CellTypes.dataCell) {
              const index = rows.findIndex((row) => row.id === rowId);

              if (hasFooter && index + 1 > rows.length - 1) {
                newCellType = CellTypes.footerCell;
                newRowIndex = -1;
                newGroupId = null;
                newrowId = null;
                break;
              }

              if (
                isGroupView &&
                (rows[index].original as TDataWithGroupId<TData>).editableTableGroupId !==
                  (rows[index + 1]?.original as TDataWithGroupId<TData>)?.editableTableGroupId
              ) {
                newGroupId = (rows[index + 1].original as TDataWithGroupId<TData>)
                  .editableTableGroupId;
                newrowId = null;
                newCellType = CellTypes.groupHeaderCell;
                newRowIndex = -1;
              } else {
                newrowId = rows[Math.min(rows.length - 1, index + 1)].id;
                newRowIndex = Math.min(rows.length - 1, index + 1);
              }
            } else if (hasFooter && cellType === CellTypes.footerCell) {
              return;
            }
            break;
          case 'ArrowLeft': {
            if (cellType === CellTypes.groupHeaderCell) {
              if (allowBulkGroupSelection && columnIndex > 0) {
                newColumnIndex = 0;
              }
            } else {
              newColumnIndex = Math.max(0, newColumnIndex - 1);
            }
            break;
          }
          case 'ArrowRight': {
            if (cellType === CellTypes.groupHeaderCell) {
              if (allowBulkGroupSelection && columnIndex === 0) {
                newColumnIndex = 1;
              }
            } else {
              newColumnIndex = Math.min(tableHeaders.length - 1, newColumnIndex + 1);
            }
            break;
          }
        }

        setFocusedCell({
          rowId: newrowId,
          columnIndex: newColumnIndex,
          cellType: newCellType,
          groupId: newGroupId,
          rowIndex: newRowIndex,
        });
      }

      if (e.key === 'Enter' && cellType === CellTypes.headerCell) {
        const header = tableHeaders[columnIndex]?.column;
        if (header?.getCanSort()) {
          header.toggleSorting();
        }
      }
    };

    const table = tableRef.current;

    table?.addEventListener('keydown', handleKeyDown);
    return () => table?.removeEventListener('keydown', handleKeyDown);
  }, [focusedCell, props.table]);

  return (
    <FocusDispatchContext.Provider value={updateData}>
      <FocusDataContext.Provider value={focusedCell}>{props.children}</FocusDataContext.Provider>
    </FocusDispatchContext.Provider>
  );
}

export function useFocusHandler({
  cellType,
  columnIndex,
  groupId,
  rowId,
  rowIndex,
}: UseFocusHandlerProps) {
  const focusedCell = useContext(FocusDataContext);
  const setFocusedCell = useContext(FocusDispatchContext);
  const { allowBulkGroupSelection, tableRef, table, serverPagination } = useTableContext();

  // TanStack pagination state
  const filteredRows = table?.getRowModel().rows ?? [];
  const localRowIndex = filteredRows.findIndex((row) => row.id === rowId);

  // Only this cell's boolean — primitive, stable comparison
  const isCurrentCellFocused = useMemo(() => {
    const isFocused =
      tableRef.current?.contains(document.activeElement) && focusedCell.cellType === cellType;

    switch (cellType) {
      case CellTypes.headerCell:
      case CellTypes.footerCell:
        return isFocused && focusedCell.columnIndex === columnIndex;
      case CellTypes.groupHeaderCell: {
        let isGroup = focusedCell.groupId === groupId;
        if (allowBulkGroupSelection) {
          isGroup &&=
            columnIndex === 0 ? focusedCell.columnIndex === 0 : focusedCell.columnIndex > 0;
        }
        return isFocused && isGroup;
      }
      case CellTypes.dataCell:
        return (
          isFocused &&
          focusedCell.rowIndex === localRowIndex &&
          focusedCell.columnIndex === columnIndex
        );
      default:
        return false;
    }
  }, [focusedCell, columnIndex, cellType, groupId, localRowIndex, allowBulkGroupSelection]);

  const focusCurrentCell = useCallback(() => {
    setFocusedCell?.({ columnIndex, rowId, cellType, groupId, rowIndex });
  }, [columnIndex, rowId, cellType, groupId, rowIndex]);

  const resetFocus = useCallback(() => {
    setFocusedCell?.({
      cellType: CellTypes.headerCell,
      rowId: null,
      groupId: null,
      rowIndex: -1,
      columnIndex: 0,
    });
  }, []);

  return {
    isCurrentCellFocused,
    focusCurrentCell,
    resetFocus,
  };
}