import React, { useEffect, useCallback, useMemo } from 'react';
import { CaretUpDownIcon } from '@trail-ui/icons';
import { ArrowDownIcon, ArrowUpIcon } from '@trail-ui/icons';
import type { SortDirection } from '@tanstack/react-table';
import { Spinner } from '../spinner';
import { Checkbox } from '../checkbox';

export interface HeadersData<DataType> {
  id: string;
  accessorKey: keyof DataType | '';
  label: string;
  isSortable?: boolean;
  width?: number | string;
  columnType?: 'th' | 'td';
  headerRenderFn?: () => React.ReactNode;
  cellRenderFn?: (data: DataType) => React.ReactNode;
}

export type NormalRow<DataType> = DataType;

export type GroupedRow<DataType> = {
  groupId: string;
  groupName: string;
  items: NormalRow<DataType>[];
};

interface BaseProps<DataType> {
  tableName: string;
  isLoading: boolean;
  headers: HeadersData<DataType>[];
  maxHeight?: string;
  selection?: {
    selectedIds: Set<string>;
    retainSelectionAccrossPages: boolean;
    setSelectedIds: React.Dispatch<React.SetStateAction<Set<string>>>;
    idSelector: (data: DataType) => string;
    checkBoxLabel?: (data: DataType) => string;
    hideSelectAll?: boolean;
    isRowSelectionDisabled?: (data: DataType) => boolean;
    onSelectAll?: () => void;
  };
  sort?: {
    sortBy: string;
    sortDirection: false | SortDirection;
    onSortClick: (sortby: string, sortDirection: false | SortDirection) => void;
  };
}

export interface NormalTableProps<DataType> extends BaseProps<DataType> {
  tableType: 'normal';
  data: NormalRow<DataType>[];
}

export interface GroupedTableProps<DataType> extends BaseProps<DataType> {
  tableType: 'grouped';
  data: GroupedRow<DataType>[];
}

export const sortOptions = {
  asc: <ArrowUpIcon aria-hidden className="h-4 w-4 text-purple-600" />,
  desc: <ArrowDownIcon aria-hidden className="h-4 w-4 text-purple-600" />,
};

export const sortOrder = [false, 'asc', 'desc'] as const;

interface SelectAllCheckboxProps<DataType> {
  tableName: string;
  data: DataType[];
  selectedIds: Set<string>;
  onToggle: (select: boolean) => void;
  idSelector: (data: DataType) => string;
  selectionCheckFn?: (data: DataType) => boolean;
}

function SelectAllCheckboxComponent<DataType>({
  tableName,
  data,
  selectedIds,
  onToggle,
  idSelector,
  selectionCheckFn,
}: SelectAllCheckboxProps<DataType>) {
  const allSelected =
    selectedIds.size > 0 &&
    data.filter((d) => !selectionCheckFn?.(d)).every((d) => selectedIds.has(idSelector(d)));
  const someSelected = data
    .filter((d) => !selectionCheckFn?.(d))
    .some((d) => selectedIds.has(idSelector(d)));

  return (
    <div className="flex justify-center">
      <Checkbox
        aria-label={`Select All ${tableName}`}
        isSelected={allSelected}
        isIndeterminate={!allSelected && someSelected}
        onChange={() => onToggle(!allSelected)}
      />
    </div>
  );
}

// ✅ Properly typed memo wrapper
export const SelectAllCheckbox = React.memo(SelectAllCheckboxComponent) as <DataType>(
  props: SelectAllCheckboxProps<DataType>,
) => React.ReactElement;

type TableProps<DataType> = NormalTableProps<DataType> | GroupedTableProps<DataType>;

export default function CustomTable<DataType>(props: TableProps<DataType>) {
  const toggleRowSelection = useCallback(
    (id: string) => {
      props.selection?.setSelectedIds((prev) => {
        const updated = new Set(prev);
        if (updated.has(id)) {
          updated.delete(id);
        } else {
          updated.add(id);
        }
        return updated;
      });
    },
    [props.selection],
  );

  const toggleAllSelection = useCallback(
    (select: boolean) => {
      if (props.selection?.onSelectAll) {
        props.selection?.onSelectAll();
      } else {
        props.selection?.setSelectedIds((prev) => {
          const updated = new Set<string>(prev);
          // biome-ignore lint/complexity/noForEach: <explanation>
          (props.data as NormalRow<DataType>[]).forEach((item) => {
            if (!props.selection?.isRowSelectionDisabled?.(item) && props.selection?.idSelector) {
              if (select) {
                updated.add(props.selection?.idSelector(item));
              } else {
                updated.delete(props.selection?.idSelector(item));
              }
            }
          });
          return updated;
        });
      }
    },
    [props.data, props.selection],
  );

  const toggleGroupSelection = useCallback(
    (items: NormalRow<DataType>[], select: boolean) => {
      props.selection?.setSelectedIds((prev) => {
        const updated = new Set(prev);
        // biome-ignore lint/complexity/noForEach: <explanation>
        items.forEach((item) => {
          if (props.selection?.idSelector) {
            const id = props.selection?.idSelector(item);
            if (select) {
              updated.add(id);
            } else {
              updated.delete(id);
            }
          }
        });
        return updated;
      });
    },
    [props.selection],
  );

  const handleSortChange = useCallback(
    (sortBy: string) => {
      if (!props.sort) {
        return;
      }
      const index = props.sort.sortBy === sortBy ? sortOrder.indexOf(props.sort.sortDirection) : 0;
      const nextDirection = sortOrder[(index + 1) % sortOrder.length];
      props.sort.onSortClick(sortBy, nextDirection);
    },
    [props.sort],
  );

  const renderRow = useCallback(
    (item: NormalRow<DataType>, index: number) => {
      const isChecked = props.selection?.selectedIds.has(props.selection.idSelector(item));

      return (
        <tr
          key={props.selection?.idSelector(item) || index}
          className="border-b border-b-neutral-200 last:border-b-0"
        >
          {props.selection && (
            <td className={`${isChecked ? 'bg-purple-50' : 'bg-neutral-50'} px-2 py-2 text-sm`}>
              <div className="flex justify-center">
                <Checkbox
                  aria-label={
                    props.selection?.checkBoxLabel
                      ? props.selection.checkBoxLabel(item)
                      : `Select ${props.selection.idSelector(item)}`
                  }
                  isSelected={isChecked}
                  onChange={() =>
                    props.selection?.idSelector &&
                    toggleRowSelection(props.selection.idSelector(item))
                  }
                  {...(props.selection.isRowSelectionDisabled && {
                    isDisabled: props.selection.isRowSelectionDisabled(item),
                  })}
                />
              </div>
            </td>
          )}
          {props.headers.map((column) => {
            const Element = column.columnType || 'td';
            return (
              <Element
                key={column.id}
                headers={column.id}
                className={`${isChecked ? 'bg-purple-50' : 'bg-neutral-50'} px-2 py-2 text-sm`}
                style={{
                  minWidth: column.width,
                  maxWidth: column.width,
                  width: column.width,
                }}
              >
                {column.cellRenderFn
                  ? column.cellRenderFn(item)
                  : column.accessorKey && (item[column.accessorKey] as React.ReactNode)}
              </Element>
            );
          })}
        </tr>
      );
    },
    [props.selection, props.headers, toggleRowSelection],
  );

  const renderGroup = useCallback(
    (group: GroupedRow<DataType>) => {
      const allSelected = group.items.every((item) =>
        props.selection?.selectedIds.has(props.selection.idSelector(item)),
      );
      const someSelected = group.items.some((item) =>
        props.selection?.selectedIds.has(props.selection.idSelector(item)),
      );

      return (
        <React.Fragment key={group.groupId}>
          <tr>
            <td className="border-b border-r bg-neutral-50 p-0">
              <div className="flex justify-center">
                <Checkbox
                  aria-label={`Select group ${group.groupName}`}
                  isSelected={allSelected}
                  isIndeterminate={!allSelected && someSelected}
                  onChange={() => toggleGroupSelection(group.items, !allSelected)}
                />
              </div>
            </td>
            <th colSpan={props.headers.length} className="border-b bg-neutral-50 text-left">
              <span className="text-sm font-semibold">{group.groupName}</span>{' '}
              <span className="text-sm font-normal text-neutral-600">
                ({group.items.length} {group.items.length === 1 ? 'issue' : 'issues'})
              </span>
            </th>
          </tr>
          {group.items.map(renderRow)}
        </React.Fragment>
      );
    },
    [props.selection, props.headers.length, toggleGroupSelection, renderRow],
  );

  const columnCells = useMemo(() => {
    return props.headers.map((column) => (
      <th
        key={column.id}
        id={column.id}
        scope="col"
        className="h-10 border-r border-neutral-200 bg-neutral-100 p-2 text-left text-sm font-semibold last:border-r-0"
        style={{
          minWidth: column.width === 100 ? 'max-content' : column.width,
          maxWidth: column.width === 100 ? 'max-content' : column.width,
          width: column.width,
        }}
        {...(column.isSortable
          ? {
              'aria-sort':
                props.sort?.sortBy === column.id
                  ? props.sort?.sortDirection === 'asc'
                    ? 'ascending'
                    : 'descending'
                  : 'none',
            }
          : {})}
      >
        {column.isSortable ? (
          <button
            type="button"
            onClick={() => handleSortChange(column.id)}
            className="flex w-full items-center justify-between gap-4 px-1 outline-purple-600"
          >
            {column.headerRenderFn ? (
              column.headerRenderFn()
            ) : (
              <span className="text-left">
                {column.label} <span className="sr-only"> Sortable</span>
              </span>
            )}
            {props.sort?.sortBy === column.id && props.sort.sortDirection !== false ? (
              sortOptions[props.sort.sortDirection]
            ) : (
              <CaretUpDownIcon className="h-4 w-4 text-neutral-600" />
            )}
          </button>
        ) : column.headerRenderFn ? (
          column.headerRenderFn()
        ) : (
          column.label
        )}
      </th>
    ));
  }, [props.headers, props.sort, handleSortChange]);

  useEffect(() => {
    if (!props.selection?.retainSelectionAccrossPages) {
      props.selection?.setSelectedIds(new Set());
    }
  }, [props.data]);

  return (
    <div className="rounded">
      <div className="h-full overflow-auto rounded border border-neutral-200">
        <div className={props.maxHeight ?? 'max-h-[60vh]'}>
          <table
            aria-label={props.tableName ?? ''}
            className="w-full table-auto border-collapse bg-white"
          >
            {props.headers.some((header) => header.isSortable) && (
              <caption className="sr-only">column headers with buttons are sortable</caption>
            )}
            <thead className="sticky top-0 z-[1] before:absolute before:bottom-0 before:w-full before:border-b before:border-b-neutral-200">
              <tr className="bg-gray-200">
                {props.selection &&
                props.tableType === 'normal' &&
                !props.selection.hideSelectAll ? (
                  <td className="h-10 w-10 border-r border-neutral-200 bg-neutral-100 p-2 text-sm font-semibold">
                    <SelectAllCheckbox
                      tableName={props.tableName}
                      data={props.data as NormalRow<DataType>[]}
                      selectedIds={props.selection.selectedIds}
                      onToggle={toggleAllSelection}
                      idSelector={props.selection.idSelector}
                      selectionCheckFn={props.selection.isRowSelectionDisabled}
                    />
                  </td>
                ) : (
                  props.selection && (
                    <td className="h-10 w-10 border-r border-neutral-200 bg-neutral-100 p-2" />
                  )
                )}
                {columnCells}
              </tr>
            </thead>
            <tbody>
              {props.isLoading ? (
                <tr className="h-[300px]">
                  <td colSpan={props.headers.length + 1}>
                    <div className="flex flex-col items-center justify-center gap-1">
                      <Spinner className="h-8 w-8 animate-spin" />
                      <p className="pt-1 text-neutral-700">Loading...</p>
                    </div>
                  </td>
                </tr>
              ) : !props.data?.length ? (
                <tr className="h-[300px]">
                  <td colSpan={props.headers.length + 1}>
                    <div className="flex flex-col items-center justify-center gap-1">
                      <p role="alert" aria-live="polite" className="text-neutral-700">
                        No data
                      </p>
                    </div>
                  </td>
                </tr>
              ) : (
                props.data.map((item, index) =>
                  props.tableType === 'normal'
                    ? renderRow(item as NormalRow<DataType>, index)
                    : renderGroup(item as GroupedRow<DataType>),
                )
              )}
            </tbody>
          </table>
        </div>
      </div>
    </div>
  );
}