import { useRef } from "react";
import { ISubRowColumnProps } from "./ISubRowColumnProps";

export const useSubRows = <TElement, TSubRow, TKey, TError>(
  elements: TElement[],
  createSubRows: (element: TElement) => TSubRow[],
  getKey: (element: TElement) => TKey,
  areElementsEqual: (a: TElement, b: TElement) => boolean,
  columnProps: ISubRowColumnProps<TElement, TSubRow, TError>[]
) => {
  const prevElements = useRef<TElement[] | null>(null);
  const prevGetKeyRef = useRef<((element: TElement) => TKey) | null>(null);
  const indexByKey = useRef<Map<
    TKey,
    { rowsStart: number; numRows: number }
  > | null>(null);

  const rowsRef = useRef<{ element: TElement; subRow: TSubRow }[] | null>(null);
  const errorsRef = useRef<(TError | null)[][] | null>(null);

  let needFullReindex = true;
  if (
    indexByKey.current &&
    rowsRef.current &&
    prevGetKeyRef.current === getKey &&
    errorsRef.current
  ) {
    if (prevElements.current !== elements) {
      if (prevElements.current?.length === elements.length) {
        needFullReindex = false;
        const changedIndices = prevElements.current
          .map((_, i) => i)
          .filter(
            (i) => !areElementsEqual(prevElements.current![i], elements[i])
          );
        const newRows = [...rowsRef.current];
        const newErrors = [...errorsRef.current];
        for (const i of changedIndices) {
          const newKey = getKey(elements[i]);
          if (newKey !== getKey(prevElements.current[i])) {
            needFullReindex = true;
            break;
          }
          const data = indexByKey.current?.get(newKey);
          if (!data) {
            needFullReindex = true;
            break;
          }
          const updatedSubRows = createSubRows(elements[i]);
          if (updatedSubRows.length !== data.numRows) {
            needFullReindex = true;
            break;
          }
          updatedSubRows.forEach((subRow, subRowIndex) => {
            newRows[data.rowsStart + subRowIndex] = {
              element: elements[i],
              subRow,
            };
            const errors = columnProps.map(
              (props) => props.validate?.(elements[i], subRow) ?? null
            );
            newErrors[data.rowsStart + subRowIndex] = errors;
          });
        }
        rowsRef.current = newRows;
        errorsRef.current = newErrors;
      }
    }
  }

  if (needFullReindex) {
    const subRowGroups = elements.map(createSubRows);
    let rowIndex = 0;
    indexByKey.current = new Map(
      subRowGroups.map((subRowGroup, index) => {
        const ret = [
          getKey(elements[index]),
          {
            elementIndex: index,
            rowsStart: rowIndex,
            numRows: subRowGroup.length,
          },
        ] as const;
        rowIndex += subRowGroup.length;
        return ret;
      })
    );
    rowsRef.current = subRowGroups.flatMap((subRows, index) =>
      subRows.map((subRow) => ({
        subRow,
        element: elements[index],
      }))
    );

    errorsRef.current = subRowGroups
      .map((subRowGroup, index) =>
        subRowGroup.map((subRow) =>
          columnProps.map(
            (props) => props.validate?.(elements[index], subRow) ?? null
          )
        )
      )
      .flat();
  }

  prevGetKeyRef.current = getKey;
  prevElements.current = elements;

  return {
    rows: rowsRef.current!,
    errors: errorsRef.current!,
  };
};