'use client'

import {
  createContext,
  useContext,
  useLayoutEffect,
  useRef,
  useState,
  type ComponentProps,
  type ReactNode,
} from 'react'
import { Dialog } from '@base-ui/react/dialog'

import type { GDSComponentProps } from '@graphprotocol/gds-css'
import { ArrowLeftIcon, XIcon } from '@graphprotocol/gds-react/icons'

import { useControlled, useCSSPropsPolyfill, useGDS } from '../hooks/index.ts'
import { cn, getCSSPropsAttributes } from '../utils/index.ts'
import { Portal } from './base/Portal.tsx'
import type { RenderProp } from './base/Render.tsx'
import { Transition } from './base/Transition.tsx'
import { Button } from './Button.tsx'
import { ButtonGroup } from './ButtonGroup.tsx'
import { Divider } from './Divider.tsx'
import { ModalMeta } from './Modal.meta.ts'
import { Stepper, type StepperProps } from './Stepper.tsx'

const ModalContext = createContext<{
  open: boolean
  setOpen: (open: boolean) => void
  dismissible: boolean
  steps: (string | number)[]
  currentStep: string | number
  setCurrentStep: (step: string | number) => void
  currentStepIndex: number | null
  currentStepDirection: 'left' | 'right'
} | null>(null)

export interface ModalProps<
  Steps extends number | readonly string[] = readonly string[],
  Step extends Steps extends readonly string[] ? Steps[number] : number =
    Steps extends readonly string[] ? Steps[number] : number,
>
  extends Omit<ComponentProps<'div'>, 'children'>, GDSComponentProps<typeof ModalMeta> {
  /**
   * Element that opens the modal when clicked, or a function that returns such an element (where
   * you can fully customize the interaction that opens the modal, e.g. `trigger={(renderProps, {
   * openModal }) => <button {...renderProps} onMouseEnter={openModal} />}`).
   */
  trigger?:
    | RenderProp<{
        open: boolean
        setOpen: (open: boolean) => void
        openModal: () => void
        currentStep: Step
        setCurrentStep: (step: Step) => void
      }>
    | undefined
  open?: boolean | undefined
  /** Defaults to `false` if `trigger` is provided, or `true` otherwise. */
  defaultOpen?: boolean | undefined
  onOpenChange?: ((open: boolean) => void) | undefined
  /**
   * Whether the user should be able to close the modal on their own, either by clicking outside of
   * it, by pressing Escape, or by clicking the close button that gets rendered in `Modal.Header`
   * when `dismissible` is `true` (though it can be replaced or removed altogether with the
   * `buttonAfter` prop). Defaults to `true` if either `trigger` or `onOpenChange` is provided, or
   * `false` otherwise.
   */
  dismissible?: boolean | undefined
  /**
   * One of `dialog` or `alertdialog`. Defaults to `dialog` if `dismissible` is `true`, or
   * `alertdialog` if it's `false`.
   */
  role?: 'dialog' | 'alertdialog' | undefined
  steps?: Steps | undefined
  currentStep?: Step | undefined
  onCurrentStepChange?: ((step: Step) => void) | undefined
  children?:
    | ReactNode
    | ((state: {
        open: boolean
        setOpen: (open: boolean) => void
        closeModal: () => void
        currentStep: Step
        setCurrentStep: (step: Step) => void
      }) => ReactNode)
    | undefined
}

export const ModalRoot = <
  Steps extends number | readonly string[],
  Step extends Steps extends readonly string[] ? Steps[number] : number,
>({
  ref: passedRef,
  variant,
  size,
  trigger,
  open: controlledOpen,
  defaultOpen,
  onOpenChange,
  dismissible: passedDismissible,
  role,
  steps: passedSteps = 1 as Steps,
  currentStep: controlledCurrentStep,
  onCurrentStepChange,
  className,
  style,
  children,
  ...props
}: ModalProps<Steps, Step>) => {
  const { dirProps } = useGDS()

  const [cssPropsPolyfillRef, cssPropsPolyfillAttributes] = useCSSPropsPolyfill(ModalMeta, {
    variant,
    size,
  })
  const cssPropsAttributes = getCSSPropsAttributes(ModalMeta, { variant, size }, style)
  const cssPropsClassName = cn(`
    u:default-size-medium
    u:default-variant-dialog
    u:max-xs:default-variant-drawer
    u:max-2xs:prop-size-small
  `)

  const hasTrigger = trigger !== undefined

  const [open, setOpen] = useControlled(controlledOpen, defaultOpen ?? !hasTrigger, onOpenChange)

  // TODO: Allow swiping down to close when `variant` is `drawer` and `dismissible` is `true`
  const dismissible = passedDismissible ?? (hasTrigger || onOpenChange !== undefined)

  const steps = (
    typeof passedSteps === 'number'
      ? Array.from({ length: Math.max(1, passedSteps) }, (_, i) => i)
      : passedSteps.length === 0
        ? [0]
        : passedSteps
  ) as Step[]
  const [currentStep, setCurrentStep] = useControlled(
    controlledCurrentStep,
    steps[0]!,
    onCurrentStepChange,
  )
  let currentStepIndex: number | null = steps.indexOf(currentStep)
  currentStepIndex = currentStepIndex === -1 ? null : currentStepIndex
  const previousStepIndex = useRef(currentStepIndex)
  const currentStepDirection = useRef<'left' | 'right'>('left')
  if (currentStepIndex !== previousStepIndex.current) {
    currentStepDirection.current =
      (currentStepIndex ?? steps.length) > (previousStepIndex.current ?? steps.length)
        ? 'left'
        : 'right'
    previousStepIndex.current = currentStepIndex
  }

  return (
    <ModalContext.Provider
      value={{
        open,
        setOpen,
        dismissible,
        steps,
        currentStep,
        setCurrentStep: (step: string | number) => {
          if (typeof step === 'number' && steps[step]) {
            setCurrentStep(steps[step])
          } else if (steps.includes(step as Step)) {
            setCurrentStep(step as Step)
          }
        },
        currentStepIndex,
        currentStepDirection: currentStepDirection.current,
      }}
    >
      <Dialog.Root
        open={open}
        onOpenChange={(newOpen) => {
          if (!newOpen && !dismissible) return
          setOpen(newOpen)
        }}
        disablePointerDismissal={!dismissible}
      >
        {hasTrigger ? (
          <Dialog.Trigger
            render={
              typeof trigger !== 'function'
                ? trigger
                : (renderProps) => {
                    // Don't pass event handlers and other potentially too specific props (e.g. `type`)
                    const filteredRenderProps = Object.fromEntries(
                      Object.entries(renderProps).filter(
                        ([prop]) =>
                          prop === 'ref' ||
                          prop === 'id' ||
                          prop === 'className' ||
                          prop.startsWith('aria-') ||
                          prop.startsWith('data-'),
                      ),
                    )
                    return trigger(filteredRenderProps, {
                      open,
                      setOpen,
                      openModal: () => setOpen(true),
                      currentStep,
                      setCurrentStep,
                    })
                  }
            }
          />
        ) : null}
        {/* Ensure the CSS props polyfill works even when the modal is not mounted */}
        <Portal>
          <div
            ref={cssPropsPolyfillRef}
            className={cn('gds-modal i:invisible', cssPropsClassName, className)}
            {...cssPropsAttributes}
            {...cssPropsPolyfillAttributes}
            {...dirProps}
            {...props}
          />
        </Portal>
        <Dialog.Portal>
          <div
            ref={passedRef}
            /**
             * Establishing the root element as a size query container so that consumers can use
             * `w-*` and `h-*` utilities on it to set the size of the nested `Dialog.Popup`, thanks
             * to `cqi` & `cqb` units as well as some (probably too clever) math tricks. For
             * instance, if a width is set on the root, `max(100cqi,100vw-100cqi*9999)` (used to
             * calculate the popup's `max-width`) returns that width (100cqi, since it's guaranteed
             * to be larger than 100vw minus 9999 times that value), but if no width is set on the
             * root, then it returns 100vw because 100cqi, even 9999 times, is 0.
             */
            className={cn(
              `gds-modal @container-[size] pointer-events-none absolute top-0 left-0 u:text-16 u:text-muted i:invisible i:*:visible`,
              cssPropsClassName,
              className,
            )}
            {...cssPropsAttributes}
            {...cssPropsPolyfillAttributes}
            {...dirProps}
            {...props}
          >
            <Dialog.Backdrop
              className={`
                fixed inset-0 bg-backdrop transition
                data-ending-style:opacity-0
                data-open:pointer-events-previous
                data-starting-style:opacity-0
                @prop-variant-dialog/modal:duration-200
                @prop-variant-drawer/modal:duration-400
              `}
            />
            <Dialog.Popup
              role={role ?? (dismissible ? 'dialog' : 'alertdialog')}
              className={`
                fixed inset-0 m-auto h-max
                max-h-[min(max(100cqb,100vh-100cqb*9999),var(--strict-max-height))]
                min-h-[min(100cqb,var(--strict-max-height))]
                max-w-[min(max(100cqi,100vw-100cqi*9999),var(--strict-max-width))]
                min-w-[min(100cqi,var(--strict-max-width))]
                origin-bottom rounded-t-(--gds-modal-radius-top) rounded-b-(--gds-modal-radius-bottom)
                border border-muted bg-muted outline-0 transition
                [--strict-max-height:max(100vh-2*var(--gds-modal-margin-y))]
                [--strict-max-width:calc(100vw-2*var(--gds-modal-margin-x))]
                data-open:pointer-events-previous
                @prop-variant-dialog/modal:duration-200
                @prop-variant-dialog/modal:data-ending-style:scale-[97%]
                @prop-variant-dialog/modal:data-ending-style:opacity-0
                @prop-variant-dialog/modal:data-starting-style:scale-[97%]
                @prop-variant-dialog/modal:data-starting-style:opacity-0
                @prop-variant-drawer/modal:mb-0
                @prop-variant-drawer/modal:border-b-0
                @prop-variant-drawer/modal:duration-400
                @prop-variant-drawer/modal:[--strict-max-height:max(100vh-var(--gds-modal-margin-y))]
                @prop-variant-drawer/modal:data-ending-style:translate-y-full
                @prop-variant-drawer/modal:data-starting-style:translate-y-full
                @prop-size-small/modal:w-80
                @prop-size-medium/modal:w-130
                @prop-size-large/modal:w-160
                i:@prop-variant-drawer/modal:max-xs:w-full
              `}
            >
              <div
                className={`
                  max-h-[inherit] min-h-[inherit] overflow-x-clip overflow-y-auto gradient-mask-y px-(--gds-modal-padding-x) py-(--gds-modal-padding-y)
                  has-nested/modal-slot:flex
                  has-nested/modal-slot:flex-col
                  has-nested/modal-slot:overflow-y-clip
                  has-nested/modal-slot:p-0
                `}
              >
                {typeof children === 'function'
                  ? children({
                      open,
                      setOpen,
                      closeModal: () => setOpen(false),
                      currentStep,
                      setCurrentStep,
                    })
                  : children}
              </div>
            </Dialog.Popup>
          </div>
        </Dialog.Portal>
      </Dialog.Root>
    </ModalContext.Provider>
  )
}
ModalRoot.displayName = 'Modal'

export interface ModalHeaderProps extends Omit<ComponentProps<'div'>, 'title' | 'children'> {
  title: ReactNode
  /** Defaults to a back button if there is a previous step, or `null` otherwise. */
  buttonBefore?: ReactNode | undefined
  /** Defaults to a close button if the modal is `dismissible`, or `null` otherwise. */
  buttonAfter?: ReactNode | undefined
  /**
   * Props to pass to the `Stepper` component, or `null` to render a `Divider` instead. Defaults to
   * match the `steps` and `currentStep` props of the parent `Modal` (except if there is only one
   * step or if the `currentStep` is not among the `steps`, in which case a divider is rendered
   * instead of the stepper).
   */
  stepper?:
    | {
        count: NonNullable<StepperProps['count']>
        currentIndex: StepperProps['currentIndex']
      }
    | null
    | undefined
}

export const ModalHeader = ({
  title,
  buttonBefore: passedButtonBefore,
  buttonAfter: passedButtonAfter,
  stepper: passedStepper,
  className,
  ...props
}: ModalHeaderProps) => {
  const modalContext = useContext(ModalContext)
  const stepCount = modalContext?.steps.length ?? 0
  const currentStep = modalContext?.currentStep ?? 0
  const currentStepIndex = modalContext?.currentStepIndex ?? null
  const currentStepDirection = modalContext?.currentStepDirection ?? 'left'
  const stepper =
    passedStepper !== undefined
      ? passedStepper && passedStepper.count >= 1
        ? passedStepper
        : null
      : stepCount > 1 && currentStepIndex !== null
        ? { count: stepCount, currentIndex: currentStepIndex }
        : null

  const buttonTransitionClassName = cn(`
    min-w-5 shrink-0
    **:button:default-variant-naked
    style-nested/transition-content:items-center
    style-nested/transition-content:justify-center
    @prop-size-small/modal:**:button:default-size-small
    @prop-size-medium/modal:**:button:default-size-medium
    @prop-size-large/modal:**:button:default-size-medium
  `)
  const buttonBefore: ReactNode =
    passedButtonBefore !== undefined ? (
      passedButtonBefore
    ) : modalContext && currentStepIndex !== null && currentStepIndex > 0 ? (
      <Button onClick={() => modalContext.setCurrentStep(currentStepIndex - 1)}>
        <ArrowLeftIcon alt="Back" />
      </Button>
    ) : null
  const buttonAfter: ReactNode =
    passedButtonAfter !== undefined ? (
      passedButtonAfter
    ) : modalContext && modalContext.dismissible ? (
      <Dialog.Close
        render={
          <Button tooltip={null}>
            <XIcon alt="Close" />
          </Button>
        }
      />
    ) : null

  return (
    <div className={cn('nested/modal-slot shrink-0', className)} {...props}>
      <div
        className={`
          flex items-center gap-4
          @prop-size-small/modal:h-12
          @prop-size-small/modal:px-4
          @prop-size-medium/modal:h-16
          @prop-size-medium/modal:px-6
          @prop-size-large/modal:h-16
          @prop-size-large/modal:px-6
        `}
      >
        <Transition
          layout="flex-col"
          mode="exit-enter"
          animateSize={false}
          className={buttonTransitionClassName}
        >
          {buttonBefore}
        </Transition>
        <Dialog.Title
          className={`
            flex-1 text-center font-medium text-default
            @prop-size-small/modal:text-14
            @prop-size-medium/modal:text-16
            @prop-size-large/modal:text-16
          `}
        >
          <Transition
            data-direction={currentStepDirection}
            className={`
              data-[direction=left]:var-[enter-translate-x=32px]
              data-[direction=left]:var-[exit-translate-x=-32px]
              data-[direction=right]:var-[enter-translate-x=-32px]
              data-[direction=right]:var-[exit-translate-x=32px]
              style-nested/transition-content:truncate
            `}
          >
            <span key={currentStep}>{title ? title : <>&nbsp;</>}</span>
          </Transition>
        </Dialog.Title>
        <Transition
          layout="flex-col"
          mode="exit-enter"
          animateSize={false}
          className={buttonTransitionClassName}
        >
          {buttonAfter}
        </Transition>
      </div>
      <Transition>
        {stepper ? <Stepper {...stepper} className="-mt-px" /> : <Divider variant="muted" />}
      </Transition>
    </div>
  )
}
ModalHeader.displayName = 'Modal.Header'

export interface ModalBodyProps extends ComponentProps<'div'> {}

export const ModalBody = ({ className, children, ...props }: ModalBodyProps) => {
  const modalContext = useContext(ModalContext)
  const currentStep = modalContext?.currentStep ?? 0

  return (
    <div className={cn('nested/modal-slot root-flex grow flex-col', className)} {...props}>
      <Transition layout="flex-col" mode="exit-enter" animateSize="clip">
        <div
          key={currentStep}
          className="max-h-full overflow-x-clip overflow-y-auto gradient-mask-y"
        >
          <div className="px-(--gds-modal-padding-x) py-(--gds-modal-padding-y)">{children}</div>
        </div>
      </Transition>
    </div>
  )
}
ModalBody.displayName = 'Modal.Body'

export interface ModalFooterProps extends ComponentProps<'div'> {
  primaryButton?: ReactNode | undefined
  secondaryButton?: ReactNode | undefined
}

export const ModalFooter = ({
  primaryButton,
  secondaryButton,
  className,
  children,
  ...props
}: ModalFooterProps) => {
  const modalContext = useContext(ModalContext)
  const currentStep = modalContext?.currentStep ?? 0
  const isEmpty = !children && !primaryButton && !secondaryButton

  return (
    <div className={cn('nested/modal-slot shrink-0', className)} {...props}>
      <Transition>{!isEmpty ? <Divider variant="muted" /> : null}</Transition>
      <Transition mode="exit-enter">
        {!isEmpty ? (
          <div
            key={currentStep}
            className="flex flex-col gap-4 px-(--gds-modal-padding-x) py-(--gds-modal-padding-y)"
          >
            {children ? <div className="text-14 text-muted">{children}</div> : null}
            {primaryButton || secondaryButton ? (
              <ButtonGroup
                fullWidth
                data-has-primary-button={Boolean(primaryButton) || undefined}
                data-has-secondary-button={Boolean(secondaryButton) || undefined}
                className={`
                  data-has-primary-button:**:button:last:default-variant-primary
                  data-has-secondary-button:**:button:default-variant-tertiary
                  @prop-size-small/modal:prop-size-small
                  @prop-size-medium/modal:prop-size-medium
                  @prop-size-large/modal:prop-size-medium
                `}
              >
                {secondaryButton}
                {primaryButton}
              </ButtonGroup>
            ) : null}
          </div>
        ) : null}
      </Transition>
    </div>
  )
}
ModalFooter.displayName = 'Modal.Footer'

const ModalStepContext = createContext<{
  version: number
  setContent: (content: ReactNode) => void
  forceRerender: () => void
} | null>(null)

export interface ModalStepProps {
  children: ReactNode
}

export const ModalStep = ({ children }: ModalStepProps) => {
  const stepContext = useContext(ModalStepContext)
  const isConsumer = stepContext !== null
  const version = useRef(0)
  const content = useRef<ReactNode>(null)
  const [, setForceRerender] = useState({})

  /**
   * If we're in a consumer and the provider's version is the same as it was after the previous
   * render, we know the current render was not triggered by the provider (since the provider always
   * increases its version when it renders), which means it's safe to re-render the provider, and
   * necessary for the updated content to be displayed.
   */
  useLayoutEffect(() => {
    if (isConsumer) {
      if (version.current === stepContext.version) {
        stepContext.forceRerender()
      } else {
        version.current = stepContext.version
      }
    }
  })

  if (isConsumer) {
    stepContext.setContent(children)
    return null
  } else {
    version.current++
    const setContent = (updatedContent: ReactNode) => {
      content.current = updatedContent
    }
    const forceRerender = () => setForceRerender({})
    return (
      <>
        <ModalStepContext.Provider value={{ version: version.current, setContent, forceRerender }}>
          {children}
        </ModalStepContext.Provider>
        <StableElement render={() => content.current} />
      </>
    )
  }
}
ModalStep.displayName = 'Modal.Step'

const StableElement = ({ render }: { render: () => ReactNode }) => render()
