import React, { useRef, useCallback, useMemo } from "react";
import * as UndoStack from "./UndoStack";

export const useUndoStack = <TRow, TKey>(
  rows: TRow[],
  getKey: (row: TRow) => TKey,
  onChange: (rows: TRow[]) => void
) => {
  const undoStackRef = useRef(UndoStack.create<TRow>());

  const rowsRef = useRef<TRow[]>(rows);
  rowsRef.current = rows;
  const indexByKey = useMemo<Map<TKey, number>>(
    () => new Map(rows.map((row, index) => [getKey(row), index])),
    [rows, getKey]
  );

  const applyChanges = useCallback(
    (changes: TRow[]) => {
      const newRows = [...rows];
      changes.forEach((change) => {
        const index = indexByKey.get(getKey(change));
        if (index === undefined) {
          return;
        }
        newRows[index] = change;
      });

      return newRows;
    },
    [rows, indexByKey, getKey]
  );

  return {
    handleUndo: useCallback(() => {
      const [newStack, changes] = UndoStack.undo(undoStackRef.current);
      undoStackRef.current = newStack;
      onChange(applyChanges(changes));
    }, [indexByKey, getKey]),

    handleRedo: useCallback(() => {
      const [newStack, changes] = UndoStack.redo(undoStackRef.current);
      undoStackRef.current = newStack;
      onChange(applyChanges(changes));
    }, [indexByKey, getKey]),

    handleChange: useCallback(
      (action: React.SetStateAction<TRow[]>) => {
        const changes =
          typeof action === "function" ? action(rowsRef.current) : action;
        const newStack = UndoStack.push(
          undoStackRef.current,
          changes,
          rowsRef.current,
          (row) => indexByKey.get(getKey(row))!
        );
        undoStackRef.current = newStack;
        onChange(applyChanges(changes));
      },
      [onChange, indexByKey, getKey]
    ),
  };
};