"use client";

import * as React from "react";
import {
  Controller,
  FormProvider,
  useFormContext,
  type ControllerProps,
  type FieldPath,
  type FieldValues,
  type SubmitHandler,
  type UseFormReturn,
} from "react-hook-form";

import { cn } from "@/lib/utils";

type FormProps<TFieldValues extends FieldValues> = Omit<
  React.FormHTMLAttributes<HTMLFormElement>,
  "onSubmit" | "children"
> & {
  form: UseFormReturn<TFieldValues>;
  onSubmit: SubmitHandler<TFieldValues>;
  children: React.ReactNode;
};

const FormItemContext = React.createContext<{ id: string } | null>(null);

type FormFieldContextValue = {
  name: string;
};

const FormFieldContext = React.createContext<FormFieldContextValue | null>(null);

const Form = <TFieldValues extends FieldValues>({
  form,
  onSubmit,
  children,
  className,
  ...props
}: FormProps<TFieldValues>) => {
  return (
    <FormProvider {...form}>
      <form
        className={cn("w-full space-y-4", className)}
        onSubmit={form.handleSubmit(onSubmit)}
        {...props}
      >
        {children}
      </form>
    </FormProvider>
  );
};

const useFormField = () => {
  const itemContext = React.useContext(FormItemContext);
  const fieldContext = React.useContext(FormFieldContext);
  const { getFieldState, formState, control } = useFormContext();

  if (!itemContext || !fieldContext) {
    throw new Error("useFormField must be used within a <FormField>.");
  }

  return {
    control,
    id: itemContext.id,
    name: fieldContext.name,
    ...getFieldState(fieldContext.name, formState),
  };
};

const FormField = <
  TFieldValues extends FieldValues,
  TName extends FieldPath<TFieldValues>,
>({
  ...props
}: Omit<ControllerProps<TFieldValues, TName>, "render"> & {
  render: ControllerProps<TFieldValues, TName>["render"];
}) => {
  return (
    <FormFieldContext.Provider value={{ name: String(props.name) }}>
      <Controller
        control={props.control}
        name={props.name}
        render={props.render}
      />
    </FormFieldContext.Provider>
  );
};

const FormItem = React.forwardRef<
  HTMLDivElement,
  React.HTMLAttributes<HTMLDivElement>
>(({ className, ...props }, ref) => {
  const id = React.useId();

  return (
    <FormItemContext.Provider value={{ id }}>
      <div ref={ref} className={cn("grid gap-2", className)} {...props} />
    </FormItemContext.Provider>
  );
});
FormItem.displayName = "FormItem";

const FormLabel = React.forwardRef<
  HTMLLabelElement,
  React.LabelHTMLAttributes<HTMLLabelElement>
>(({ className, ...props }, ref) => {
  const { error, id } = useFormField();

  return (
    <label
      ref={ref}
      htmlFor={props.htmlFor ?? id}
      className={cn("text-sm leading-none font-medium", className)}
      style={error ? { color: "var(--destructive)" } : undefined}
      {...props}
    />
  );
});
FormLabel.displayName = "FormLabel";

const getFormControlAriaDescribedBy = (fieldId: string, hasError: boolean) => {
  const descriptionId = `${fieldId}-description`;
  const messageId = `${fieldId}-message`;

  if (hasError) {
    return `${descriptionId} ${messageId}`;
  }

  return descriptionId;
};

const FormControl = React.forwardRef<
  HTMLElement,
  React.HTMLAttributes<HTMLElement>
>(({ className, children, ...props }, ref) => {
  const { id, error } = useFormField();
  const controlId = props.id ?? id;

  const control = React.Children.only(children);

  if (!React.isValidElement(control)) {
    return null;
  }

  const typedControl = control as React.ReactElement<
    React.ClassAttributes<unknown> & Record<string, unknown>
  >;

  const controlClassName = (typedControl.props as { className?: string })
    .className;

  return React.cloneElement(typedControl, {
    id: controlId,
    ref: ref,
    className: cn("relative", className, controlClassName),
    ...props,
    "aria-invalid": error ? "true" : "false",
    "aria-describedby": getFormControlAriaDescribedBy(
      controlId,
      !!error?.message,
    ),
    "aria-errormessage": error?.message ? `${controlId}-message` : undefined,
  });
});
FormControl.displayName = "FormControl";

const FormDescription = React.forwardRef<
  HTMLParagraphElement,
  React.HTMLAttributes<HTMLParagraphElement>
>(({ className, ...props }, ref) => {
  const { id } = useFormField();

  return (
    <p
      ref={ref}
      id={`${id}-description`}
      className={cn("text-xs leading-5 text-muted-foreground", className)}
      {...props}
    />
  );
});
FormDescription.displayName = "FormDescription";

const FormMessage = React.forwardRef<
  HTMLParagraphElement,
  React.HTMLAttributes<HTMLParagraphElement>
>(({ className, ...props }, ref) => {
  const { error, id } = useFormField();

  if (!error?.message) {
    return null;
  }

  return (
    <p
      ref={ref}
      id={`${id}-message`}
      className={cn("text-xs text-destructive", className)}
      role="alert"
      {...props}
    >
      {typeof error.message === "string" ? error.message : String(error.message)}
    </p>
  );
});
FormMessage.displayName = "FormMessage";

export {
  Form,
  useFormField,
  FormField,
  FormItem,
  FormLabel,
  FormControl,
  FormDescription,
  FormMessage,
};