"use client";

import React, { useEffect, forwardRef, useRef } from "react";
import { BaseComponentProps } from "../../types";

export interface ModalProps extends BaseComponentProps {
  open?: boolean;
  onClose?: () => void;
  onOpenChange?: (open: boolean) => void;
  size?: "sm" | "md" | "lg" | "xl" | "full";
  showCloseButton?: boolean;
  closeOnOverlayClick?: boolean;
  closeOnEscape?: boolean;
  role?: "dialog" | "alertdialog";
  ariaLabelledBy?: string;
  ariaDescribedBy?: string;
  animation?: boolean;
}

export const Modal = forwardRef<HTMLDivElement, ModalProps>(
  (
    {
      children,
      className = "",
      open = false,
      onClose,
      onOpenChange,
      size = "md",
      showCloseButton = true,
      closeOnOverlayClick = true,
      closeOnEscape = true,
      role = "dialog",
      ariaLabelledBy,
      ariaDescribedBy,
      animation = true,
      ...props
    },
    ref
  ) => {
    const previousFocusRef = useRef<HTMLElement | null>(null);
    const modalRef = useRef<HTMLDivElement>(null);

    const handleClose = () => {
      onClose?.();
      onOpenChange?.(false);
    };

    useEffect(() => {
      if (open) {
        // Store the previously focused element
        previousFocusRef.current = document.activeElement as HTMLElement;

        // Focus the modal content
        setTimeout(() => {
          const modalElement = modalRef.current;
          if (modalElement) {
            const focusableElement = modalElement.querySelector(
              'button:not([aria-label="Close"]), [href], input, select, textarea, [tabindex]:not([tabindex="-1"])'
            ) as HTMLElement;

            if (focusableElement) {
              focusableElement.focus();
            } else {
              modalElement.focus();
            }
          }
        }, 0);
      }

      return () => {
        // Restore focus when modal closes
        if (!open && previousFocusRef.current) {
          previousFocusRef.current.focus();
        }
      };
    }, [open]);

    useEffect(() => {
      const handleEscape = (event: KeyboardEvent) => {
        if (event.key === "Escape" && closeOnEscape && handleClose) {
          handleClose();
        }
      };

      const handleTabTrap = (event: KeyboardEvent) => {
        if (event.key === "Tab" && modalRef.current) {
          const focusableElements = modalRef.current.querySelectorAll(
            'button:not([aria-label="Close"]), [href], input, select, textarea, [tabindex]:not([tabindex="-1"])'
          );

          const firstElement = focusableElements[0] as HTMLElement;
          const lastElement = focusableElements[
            focusableElements.length - 1
          ] as HTMLElement;

          if (event.shiftKey) {
            if (document.activeElement === firstElement) {
              lastElement?.focus();
              event.preventDefault();
            }
          } else {
            if (document.activeElement === lastElement) {
              firstElement?.focus();
              event.preventDefault();
            }
          }
        }
      };

      if (open) {
        document.addEventListener("keydown", handleEscape);
        document.addEventListener("keydown", handleTabTrap);
        document.body.style.overflow = "hidden";

        // Set aria-hidden on all other elements
        const rootElements = document.querySelectorAll(
          "body > *:not([aria-live])"
        );
        rootElements.forEach((element) => {
          if (element !== modalRef.current?.closest("[data-modal-root]")) {
            element.setAttribute("aria-hidden", "true");
          }
        });
      }

      return () => {
        document.removeEventListener("keydown", handleEscape);
        document.removeEventListener("keydown", handleTabTrap);
        document.body.style.overflow = "unset";

        // Remove aria-hidden from all elements
        const hiddenElements = document.querySelectorAll(
          '[aria-hidden="true"]'
        );
        hiddenElements.forEach((element) => {
          element.removeAttribute("aria-hidden");
        });
      };
    }, [open, closeOnEscape, handleClose]);

    if (!open) return null;

    const sizeClasses = {
      sm: "max-w-md",
      md: "max-w-lg",
      lg: "max-w-2xl",
      xl: "max-w-4xl",
      full: "max-w-[95vw] max-h-[95vh]",
    };

    const handleBackdropClick = (event: React.MouseEvent) => {
      if (event.target === event.currentTarget && closeOnOverlayClick) {
        handleClose();
      }
    };

    return (
      <div
        data-modal-root
        className={`fixed inset-0 z-50 flex items-center justify-center bg-black/50 p-4 ${
          animation ? "animate-in fade-in-0 duration-300" : ""
        }`}
        onClick={handleBackdropClick}
        aria-modal="true"
        role="presentation"
      >
        <div
          ref={(node) => {
            if (typeof ref === "function") {
              ref(node);
            } else if (ref) {
              ref.current = node;
            }
            // modalRef is handled by the callback
            (
              modalRef as React.MutableRefObject<HTMLDivElement | null>
            ).current = node;
          }}
          className={`relative w-full ${
            sizeClasses[size]
          } bg-background rounded-lg shadow-lg border ${
            animation ? "animate-in zoom-in-95 duration-300" : ""
          } ${className}`}
          role={role}
          aria-labelledby={ariaLabelledBy}
          aria-describedby={ariaDescribedBy}
          tabIndex={-1}
          {...props}
        >
          {showCloseButton && (
            <button
              onClick={handleClose}
              className="absolute right-4 top-4 rounded-sm opacity-70 ring-offset-background transition-opacity hover:opacity-100 focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 z-10"
              aria-label="Close"
            >
              <svg
                className="h-4 w-4"
                fill="none"
                viewBox="0 0 24 24"
                stroke="currentColor"
                aria-hidden="true"
              >
                <path
                  strokeLinecap="round"
                  strokeLinejoin="round"
                  strokeWidth={2}
                  d="M6 18L18 6M6 6l12 12"
                />
              </svg>
            </button>
          )}
          {children}
        </div>
      </div>
    );
  }
);

Modal.displayName = "Modal";

export const ModalHeader = forwardRef<HTMLDivElement, BaseComponentProps>(
  ({ className = "", children, ...props }, ref) => {
    return (
      <div
        ref={ref}
        className={`flex flex-col space-y-1.5 text-center sm:text-left p-6 pb-4 ${className}`}
        {...props}
      >
        {children}
      </div>
    );
  }
);

ModalHeader.displayName = "ModalHeader";

export interface ModalTitleProps extends BaseComponentProps {
  as?: "h1" | "h2" | "h3" | "h4" | "h5" | "h6";
}

export const ModalTitle = forwardRef<HTMLHeadingElement, ModalTitleProps>(
  ({ className = "", as: Component = "h2", children, ...props }, ref) => {
    return (
      <Component
        ref={ref}
        className={`text-lg font-semibold leading-none tracking-tight ${className}`}
        {...props}
      >
        {children}
      </Component>
    );
  }
);

ModalTitle.displayName = "ModalTitle";

export const ModalDescription = forwardRef<
  HTMLParagraphElement,
  BaseComponentProps
>(({ className = "", children, ...props }, ref) => {
  return (
    <p
      ref={ref}
      className={`text-sm text-muted-foreground leading-relaxed ${className}`}
      {...props}
    >
      {children}
    </p>
  );
});

ModalDescription.displayName = "ModalDescription";

export const ModalContent = forwardRef<HTMLDivElement, BaseComponentProps>(
  ({ className = "", children, ...props }, ref) => {
    return (
      <div ref={ref} className={`px-6 pb-4 ${className}`} {...props}>
        {children}
      </div>
    );
  }
);

ModalContent.displayName = "ModalContent";

export interface ModalFooterProps extends BaseComponentProps {
  justify?: "start" | "end" | "center" | "between";
}

export const ModalFooter = forwardRef<HTMLDivElement, ModalFooterProps>(
  ({ className = "", justify = "end", children, ...props }, ref) => {
    const justifyClasses = {
      start: "justify-start",
      end: "justify-end",
      center: "justify-center",
      between: "justify-between",
    };

    return (
      <div
        ref={ref}
        className={`flex flex-col-reverse sm:flex-row sm:space-x-2 p-6 pt-4 border-t ${justifyClasses[justify]} ${className}`}
        {...props}
      >
        {children}
      </div>
    );
  }
);

ModalFooter.displayName = "ModalFooter";