import { Loader2 } from "lucide-react";
import { Button } from "./ui/button";
import React from "react";
import { cn } from "@/lib/utils";
import { ButtonProps } from "./ui/button";

interface ActionButtonProps extends ButtonProps {
	isLoading?: boolean;
	isSuccess?: boolean;
	error?: any;
	className?: string;
	children?: React.ReactNode; // Allow custom subcomponents as children
}

const ActionButtonDefault = React.forwardRef<
	HTMLSpanElement,
	React.HTMLAttributes<HTMLSpanElement>
>(({ className, children, ...props }, ref) => (
	<span ref={ref} className={cn("flex items-center", className)} {...props}>
		{children ?? "Submit"}
	</span>
));

const ActionButtonLoading = React.forwardRef<
	HTMLSpanElement,
	React.HTMLAttributes<HTMLSpanElement>
>(({ className, children, ...props }, ref) => (
	<ActionButtonDefault ref={ref} className={className} {...props}>
		{children ?? (
			<>
				<Loader2 className="animate-spin mr-2" /> Loading...
			</>
		)}
	</ActionButtonDefault>
));

const ActionButtonSuccess = React.forwardRef<
	HTMLSpanElement,
	React.HTMLAttributes<HTMLSpanElement>
>(({ className, children, ...props }, ref) => (
	<ActionButtonDefault ref={ref} className={className} {...props}>
		{children ?? <>Success!</>}
	</ActionButtonDefault>
));

const ActionButtonError = React.forwardRef<
	HTMLSpanElement,
	React.HTMLAttributes<HTMLSpanElement>
>(({ className, children, ...props }, ref) => (
	<ActionButtonDefault ref={ref} className={className} {...props}>
		{children ?? <>Error occurred.</>}
	</ActionButtonDefault>
));

const ActionButton = React.forwardRef<HTMLButtonElement, ActionButtonProps>(
	(
		{ variant, isLoading, isSuccess, error, className, children, ...props },
		ref,
	) => {
		let buttonContent = null;
		React.Children.forEach(children, (child) => {
			if (React.isValidElement(child)) {
				if (child.type === ActionButtonLoading && isLoading) {
					buttonContent = child;
				} else if (child.type === ActionButtonSuccess && isSuccess) {
					buttonContent = child;
				} else if (child.type === ActionButtonError && error) {
					buttonContent = child;
				} else if (
					child.type === ActionButtonDefault &&
					!isLoading &&
					!isSuccess &&
					!error
				) {
					buttonContent = child;
				}
			}
		});

		return (
			<Button
				ref={ref}
				disabled={isLoading || isSuccess || error !== undefined}
				type="submit"
				className={`w-full transition-all ease-in-out ${isSuccess ? "bg-green-800" : error ? "bg-red-800" : ""} ${className}`}
				{...props}
			>
				{buttonContent}
			</Button>
		);
	},
);

export {
	ActionButton,
	ActionButtonLoading,
	ActionButtonError,
	ActionButtonSuccess,
	ActionButtonDefault,
};