import { getIn, setIn } from 'lodash-redux-immutability';
import { useCallback, useMemo, useRef } from 'react';
import { useLocalStorage } from 'usehooks-ts';

/**
 * A hook for memorizing state values at specific paths and retrieving them for default state initialization.
 * This generalizes the pattern of storing form state values in localStorage and retrieving them.
 *
 * @param storageKey - The key to use for storing memorized values in localStorage
 * @returns An object with functions to get and set memorized values at specific paths
 */
export function useStateMemorization<T extends Record<string, any>>(
  storageKey: string,
  pathsToApply: (string | number)[][]
) {
  // Store the entire memorized state object in localStorage
  const [memorizedState, setMemorizedState] = useLocalStorage<T | null>(
    storageKey,
    null
  );

  /**
   * Retrieves a memorized value from a specific path in the state tree.
   * If no value is memorized at that path, returns the provided fallback value.
   *
   * @param path - Array of keys representing the path to the value (e.g., ['global', 'model'])
   * @param fallback - Optional fallback value if nothing is memorized at this path
   * @returns The memorized value or the fallback
   */
  const getMemorizedStateFromPath = useCallback(
    <V>(path: (string | number)[], fallback?: V): V | undefined => {
      if (!memorizedState) {
        return fallback;
      }

      const value = getIn(memorizedState, path);
      return value !== undefined ? value : fallback;
    },
    [memorizedState]
  );

  /**
   * Memorizes a value at a specific path in the state tree.
   * This will update the localStorage with the new value.
   *
   * @param path - Array of keys representing the path to store the value (e.g., ['global', 'model'])
   * @param value - The value to memorize
   */
  const memorizeStateAtPath = useCallback(
    (path: (string | number)[], value: any) => {
      setMemorizedState((prevState) => {
        const currentState = prevState || ({} as T);
        return setIn(currentState, path, value);
      });
    },
    [setMemorizedState]
  );

  const previousValuesRef = useRef<Map<string, any>>(new Map());
  const onStateChange = useCallback(
    (state: T) => {
      pathsToApply.forEach((path) => {
        const pathKey = path.join('.');
        const currentValue = getIn(state, path);
        const previousValue = previousValuesRef.current.get(pathKey);

        // Only memorize if the value has actually changed
        if (currentValue !== previousValue) {
          memorizeStateAtPath(path, currentValue);
          previousValuesRef.current.set(pathKey, currentValue);
        }
      });
    },
    [memorizeStateAtPath, pathsToApply]
  );

  /**
   * Applies memorized values to a default state object.
   * This is useful for initializing state with memorized values.
   *
   * @param defaultState - The default state object
   * @returns The default state with memorized values applied
   */
  const applyMemorizedValuesToDefaultState = useCallback(
    (defaultState: T): T => {
      if (!memorizedState) {
        return defaultState;
      }

      let result = defaultState;
      pathsToApply.forEach((path) => {
        const memorizedValue = getIn(memorizedState, path);
        if (memorizedValue !== undefined) {
          result = setIn(result, path, memorizedValue);
        }
      });

      return result;
    },
    [memorizedState, pathsToApply]
  );

  return useMemo(
    () => ({
      getMemorizedStateFromPath,
      onStateChange,
      applyMemorizedValuesToDefaultState,
    }),
    [
      getMemorizedStateFromPath,
      onStateChange,
      applyMemorizedValuesToDefaultState,
    ]
  );
}