'use client'
import * as LabelPrimitive from '@radix-ui/react-label'
import {Slot} from '@radix-ui/react-slot'
import {
  Controller,
  FormProvider,
  ControllerProps,
  SubmitHandler,
  FieldValues, useFormContext, FieldPath, UseFormReturn, ControllerRenderProps,
  ControllerFieldState, UseFormStateReturn, FieldError, SubmitErrorHandler,
} from 'react-hook-form'

import {merge} from '@/lib/utils'
import {Label} from '@/components/ui/label'

import React, {ComponentProps, createContext, ReactNode, useContext, useId} from 'react'

type FormProps<T extends FieldValues> = {
  form: UseFormReturn<T>
  onSubmit?: SubmitHandler<T>
  onError?: SubmitErrorHandler<T>
  handler?: (e?: React.BaseSyntheticEvent) => Promise<void>
} & Omit<ComponentProps<'form'>, 'onSubmit' | 'onError'>

function Form<T extends FieldValues>(rawProps: FormProps<T>) {
  const {children, onSubmit, onError, handler, ...props} = rawProps
  const form = props.form

  const handle = handler || form.handleSubmit(
    onSubmit || ((_) => {}),
    onError,
  )

  return (
    <FormProvider {...form}>
      <form
        {...props}
        onSubmit={async (event) => {
          event.preventDefault()
          event.stopPropagation()
          await handle(event)
        }}>
        {children}
      </form>
    </FormProvider>
  )
}

type FormFieldProps<
  V extends FieldValues = FieldValues,
  N extends FieldPath<V> = FieldPath<V>,
> = Omit<ControllerProps<V, N>, 'control' | 'render'> & Omit<ComponentProps<'div'>, 'children'> & {
  label?: ReactNode
  description?: ReactNode
  children: (props: {
    id: string
    field: ControllerRenderProps<V, N>
    fieldState: ControllerFieldState
    formState: UseFormStateReturn<V>
  }) => ReactNode
  classNames?: {
    label?: string
    description?: string
    message?: string
  }
}

function FormField<
  V extends FieldValues = FieldValues,
  N extends FieldPath<V> = FieldPath<V>,
>(props: FormFieldProps<V, N>) {
  const form = useFormContext<V>()
  const id = useId()
  return (
    <Controller<V, N>
      name={props.name}
      control={form.control}
      render={({field, fieldState, formState}) => (
        <div data-slot="form-field" className={merge('flex flex-col gap-2', props.className)}>

          {/* label */}
          {!!props.label
            && (
              <FormLabel id={`${id}-label`} error={fieldState.error} className={props.classNames?.label}>
                {props.label}
              </FormLabel>
            )
          }

          {/* control */}
          <Slot
            data-slot="form-control"
            aria-invalid={!!fieldState.error}
            aria-describedby={
              !!fieldState.error
                ? `${id}-description`
                : `${id}-description ${id}-message`
            }>
            {props.children({id, field, fieldState, formState})}
          </Slot>

          {/* description */}
          {!!props.description && (
            <FormDescription
              id={`${id}-description`}
              error={fieldState.error}
              className={merge(
                `text-weak`,
                props.classNames?.description,
              )}>
              {props.description}
            </FormDescription>
          )}

          {/* message */}
          {!fieldState.error ? null : (
            <FormMessage id={`${id}-message`} error={fieldState.error} className={props.classNames?.message}/>
          )}
        </div>
      )}/>
  )
}

type FormState = {
  error?: FieldError
}

function FormLabel({className, id, error, ...props}: ComponentProps<typeof LabelPrimitive.Root> & FormState) {
  return (
    <Label
      data-slot="form-label"
      data-fail={!!error}
      className={merge('data-[fail=true]:text-fail', className)}
      htmlFor={id}
      {...props}
    />
  )
}

function FormDescription({className, id, error, ...props}: ComponentProps<'p'> & FormState) {
  return (
    <p
      data-slot="form-description"
      id={`${id}-description`}
      className={merge('text-muted-foreground text-sm', className)}
      {...props}
    />
  )
}

function FormMessage({className, id, error, ...props}: ComponentProps<'p'> & FormState) {
  const body = error ? String(error?.message ?? '') : props.children

  if (!body) {
    return null
  }

  return (
    <p
      data-slot="form-message"
      id={`${id}-message`}
      className={merge('text-fail text-sm', className)}
      {...props}
    >
      {body}
    </p>
  )
}

export {
  Form,
  FormField,
  FormLabel,
  FormDescription,
  FormMessage,
}