import { useEffect, useId, useMemo, useRef, useState } from 'react';
import { CellTypes } from '../constants';
import { useFocusHandler } from '../focus-handler/focus-handler-provider';
import { useTableContext } from '../table-context/table-context-provider';
import type { TableContextType, TDataWithGroupId } from '../types';
import CellCheckbox from '../components-for-stories/checkbox';

interface GroupedRowHeaderProps {
  groupId: string;
}

export default function GroupedRowHeader<TData>(props: Readonly<GroupedRowHeaderProps>) {
  const { groupData, table, allowBulkGroupSelection } = useTableContext();

  const currentGroupData = groupData.find((d) => d.groupId === props.groupId);

  const columns = table.getHeaderGroups()[0].headers.length;

  return (
    <tr>
      {allowBulkGroupSelection && currentGroupData && (
        <GroupHeaderRowCheckbox<TData>
          groupData={currentGroupData}
          columnIndex={allowBulkGroupSelection ? 1 : 0}
        />
      )}
      {currentGroupData && (
        <GroupRowHeaderCell<TData>
          groupData={currentGroupData}
          columnIndex={allowBulkGroupSelection ? 1 : 0}
          colSpan={allowBulkGroupSelection ? columns - 1 : columns}
        />
      )}
    </tr>
  );
}

function GroupHeaderRowCheckbox<TData>(
  props: Readonly<{
    groupData: TableContextType<TData>['groupData'][0];
    columnIndex: number;
  }>,
) {
  const [isChecked, setIsChecked] = useState<boolean | 'mixed'>(false);

  const cellRef = useRef<HTMLTableCellElement>(null);
  const selectAllId = useId();

  const { table } = useTableContext();

  const { focusCurrentCell, isCurrentCellFocused } = useFocusHandler({
    cellType: CellTypes.groupHeaderCell,
    columnIndex: 0,
    groupId: props.groupData?.groupId,
    rowId: null,
    rowIndex: -1,
  });

  const currentPageRowIds = useMemo(
    () =>
      new Set(
        table
          .getRowModel()
          .rows.filter(
            (r) =>
              (r.original as TDataWithGroupId<TData>).editableTableGroupId ===
              props.groupData.groupId,
          )
          .map((r) => r.id),
      ),
    [table.getState().pagination.pageSize, table.getState().pagination.pageIndex],
  );

  function toggleRowsSelection() {
    let newValue = false;
    if (!isChecked || isChecked === 'mixed') {
      newValue = true;
    }

    currentPageRowIds.forEach((id) => {
      table.getRowModel().rowsById[id].toggleSelected(newValue);
    });
  }

  useEffect(() => {
    if (isCurrentCellFocused) {
      cellRef.current?.focus();
    }
  }, [isCurrentCellFocused]);

  useEffect(() => {
    const ids = props.groupData.ids.filter((id) => currentPageRowIds.has(id));
    const selectedRowIds = Object.keys(table.getState().rowSelection);

    if (hasAll(ids, selectedRowIds)) {
      setIsChecked(true);
    } else if (hasSome(ids, selectedRowIds)) {
      setIsChecked('mixed');
    } else {
      setIsChecked(false);
    }
  }, [
    table.getState().rowSelection,
    table.getState().pagination.pageSize,
    table.getState().pagination.pageIndex,
  ]);

  return (
    <td className="border-b border-r border-t border-neutral-200 bg-neutral-50 px-2 py-2 text-left text-sm font-normal text-neutral-900 last:border-r-0">
      <span id={selectAllId} className="hidden">
        Select all instances under
      </span>
      <div
        aria-labelledby={`${selectAllId} ${props.groupData.groupId}`}
        // biome-ignore lint/a11y/useSemanticElements: <explanation>
        role="checkbox"
        aria-checked={isChecked}
        ref={cellRef}
        className="td-content flex h-6 w-6 items-center justify-center outline-purple-600"
        onClick={(e) => {
          // If user clicked directly on <input>, don't trigger another click
          if ((e.target as HTMLElement).tagName.toLowerCase() === 'input') {
            focusCurrentCell();
            return;
          }

          const checkbox = cellRef.current?.querySelector('input[type="checkbox"]');
          (checkbox as HTMLInputElement)?.click();
          focusCurrentCell();
        }}
        onKeyDown={(e) => {
          if (e.key === 'Enter' || e.key === ' ') {
            e.preventDefault();
            const checkbox = cellRef.current?.querySelector('input[type="checkbox"]');
            if (checkbox) {
              (checkbox as HTMLInputElement).click();
            }
          }
        }}
        tabIndex={isCurrentCellFocused ? 0 : -1}
      >
        <CellCheckbox
          aria-hidden={true}
          tabIndex={-1}
          type="checkbox"
          aria-label="Select all status"
          checked={!!isChecked}
          isIndeterminate={isChecked === 'mixed'}
          isSelected={isChecked}
          onChange={toggleRowsSelection}
        />
      </div>
    </td>
  );
}

function GroupRowHeaderCell<TData>(props: {
  groupData: TableContextType<TData>['groupData'][0];
  columnIndex: number;
  colSpan: number;
}) {
  const cellRef = useRef<HTMLTableCellElement>(null);

  const { focusCurrentCell, isCurrentCellFocused } = useFocusHandler({
    cellType: CellTypes.groupHeaderCell,
    columnIndex: props.columnIndex,
    groupId: props.groupData?.groupId,
    rowId: null,
    rowIndex: -1,
  });

  useEffect(() => {
    if (isCurrentCellFocused) {
      cellRef.current?.focus();
    }
  }, [isCurrentCellFocused]);

  return (
    // biome-ignore lint/a11y/useKeyWithClickEvents: <explanation>
    <th
      id={props.groupData.groupId}
      ref={cellRef}
      // biome-ignore lint/style/noUnusedTemplateLiteral: <explanation>
      className={`border-b border-r border-t border-neutral-200 px-2 py-2 text-left text-sm font-normal text-neutral-900 last:border-r-0 focus:outline-purple-700 focus:[outline-style:solid]`}
      onClick={focusCurrentCell}
      tabIndex={isCurrentCellFocused ? 0 : -1}
      colSpan={props.colSpan}
      scope="colgroup"
    >
      <span className="overflow-hidden text-sm font-semibold">{props.groupData.groupName}</span>{' '}
      {props.groupData.description && (
        <span className="overflow-hidden text-sm text-neutral-600">
          {props.groupData.description}
        </span>
      )}
    </th>
  );
}

function hasSome<itemType>(array1: itemType[], array2: itemType[]) {
  const set2 = new Set(array2);
  return array1.some((item) => set2.has(item));
}

function hasAll<itemType>(array1: itemType[], array2: itemType[]) {
  const set2 = new Set(array2);
  return array1.every((item) => set2.has(item));
}