import type { ForwardedRef } from 'react';
import { useContext, useEffect, useState } from 'react';
import React from 'react';
import { createLeafComponent, TreeStateContext } from 'react-aria-components';
import type { Node } from 'react-stately';
import { LoaderNode } from '@react-aria/collections';

import type { Exact, SafeOmit } from '@shoptet/utils';

import type { TreeItemKey } from '../useTreeState';

export interface InheritedTreeItemLoaderDivProps extends SafeOmit<React.HTMLAttributes<HTMLDivElement>, 'onLoad'> {}

export interface TreeItemLoaderProps<V extends TreeItemKey = string> extends InheritedTreeItemLoaderDivProps {
  onLoad: (key: V) => Promise<any>;
  loading?: boolean;
}

export const TreeItemLoader = createLeafComponent(LoaderNode, function TreeItemLoading<
  T extends object,
  V extends TreeItemKey = string,
>({ onLoad, loading: loadingProp, ...rest }: TreeItemLoaderProps<V>, ref: ForwardedRef<HTMLDivElement>, item: Node<T>) {
  rest satisfies Exact<InheritedTreeItemLoaderDivProps, typeof rest>;

  const state = useContext(TreeStateContext)!;

  const [localLoading, setLocalLoading] = useState(loadingProp);
  const loading = loadingProp || localLoading;

  useEffect(() => {
    if (loading) {
      return;
    }

    const parentNode = item.parentKey && state.collection.getItem(item.parentKey);

    if (!parentNode) {
      return;
    }

    if (state.expandedKeys.has(parentNode.key)) {
      setLocalLoading(true);
      onLoad(parentNode.key as V)
        .catch(() => {
          state.setExpandedKeys(new Set([...state.expandedKeys].filter(key => key !== parentNode.key)));
        })
        .finally(() => setLocalLoading(false));
    }
  }, [state, item, loading, onLoad]);

  return <div ref={ref} {...rest} />;
});