import { mergeRefs } from '@trail-ui/hooks';
import { clsx } from '@trail-ui/shared-utils';
import type {
  DrawerReturnType,
  ModalReturnType,
  ModalSlots,
  ModalVariantProps,
  SlotsToClasses,
} from '@trail-ui/theme';
import { modal } from '@trail-ui/theme';
import { ForwardedRef, createContext, forwardRef, useEffect, useMemo, useRef } from 'react';
import type { ModalOverlayProps } from 'react-aria-components';
import { Modal as AriaModal, ModalOverlay } from 'react-aria-components';

export interface ModalProps extends ModalOverlayProps, ModalVariantProps {
  /**
   * Classes object to style the modal and its children.
   */
  classNames?: SlotsToClasses<ModalSlots>;
}

interface InternalModalContextValue {
  slots: ModalReturnType | DrawerReturnType;
  classNames?: SlotsToClasses<ModalSlots>;
  // state: ModalRenderProps['state'];
}

export const InternalModalContext = createContext<InternalModalContextValue>(
  {} as InternalModalContextValue,
);

function Modal(props: ModalProps, ref: ForwardedRef<HTMLDivElement>) {
  const {
    children,
    classNames,
    className,
    size,
    radius,
    placement,
    shadow,
    backdrop = 'opaque',
    scrollBehavior,
    ...otherProps
  } = props;

  const localRef = useRef<HTMLElement>(null);

  const slots = useMemo(
    () => modal({ size, radius, placement, shadow, scrollBehavior, backdrop }),
    [backdrop, placement, radius, scrollBehavior, shadow, size],
  );

  const baseStyles = clsx(classNames?.base, className);

  useEffect(() => {
    if (props.isOpen) {
      const childSection = localRef.current?.getElementsByTagName('section')[0];
      const labelledBy = childSection?.getAttribute('aria-labelledby');
      const role = childSection?.getAttribute('role');
      const tabIndex = childSection?.getAttribute('tabindex');
      if (labelledBy) {
        localRef.current?.setAttribute('aria-labelledby', labelledBy);
        childSection?.removeAttribute('aria-labelledby');
      }
      if (role) {
        localRef.current?.setAttribute('role', role);
        childSection?.removeAttribute('role');
      }
      if (tabIndex) {
        localRef.current?.setAttribute('tabindex', tabIndex);
        childSection?.removeAttribute('tabindex');
      }
      localRef.current?.setAttribute('aria-modal', 'true');
    }
  }, [props.isOpen]);

  return (
    <ModalOverlay {...otherProps} className={slots.backdrop({ class: classNames?.backdrop })}>
      <div className={slots.wrapper({ class: classNames?.wrapper })} data-placement={placement}>
        <InternalModalContext.Provider value={{ slots, classNames }}>
          <AriaModal ref={mergeRefs(localRef, ref)} className={slots.base({ class: baseStyles })}>
            {children}
          </AriaModal>
        </InternalModalContext.Provider>
      </div>
    </ModalOverlay>
  );
}

/**
 * A modal is an overlay element which blocks interaction with elements outside it.
 */

const _Modal = forwardRef(Modal);
export { _Modal as Modal };