web/src/components/ui/form.tsx

'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,
} from 'react-hook-form'

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

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

type FormProps<T extends FieldValues> = {
  form: UseFormReturn<T>
  onSubmit: SubmitHandler<T>
} & Omit<ComponentProps<'form'>, 'onSubmit'>

function Form<T extends FieldValues>(rawProps: FormProps<T>) {

  const {children, onSubmit, ...props} = rawProps
  const form = props.form

  return (
    <FormProvider {...form}>
      <form {...props} onSubmit={form.handleSubmit(onSubmit)}>
        {children}
      </form>
    </FormProvider>
  )
}

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

type FormFieldContext = {
  id: string
  error?: FieldError
}

const FormFieldContext = createContext<FormFieldContext | null>(null)

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('grid gap-2', props.className)}>
        <FormFieldContext value={{id: id, error: fieldState.error}}>
          {!!props.label &&
            <Label
              data-slot="form-label"
              data-error={!!fieldState.error}
              className={merge('data-[error=true]:text-destructive')}
              htmlFor={id}>
              {props.label}
            </Label>
          }

          <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>

          {!fieldState.error ? null : (
            <p
              data-slot="form-message"
              className={merge('text-destructive text-sm')}>
              {fieldState.error?.message}
            </p>
          )}
        </FormFieldContext>
      </div>
    )}
    />
  )
}

const useFormField = () => {
  const context = useContext(FormFieldContext)
  if (!context) {
    throw new Error('FormField components must be used within a FormField component')
  }
  return context
}


function FormLabel({className, ...props}: ComponentProps<typeof LabelPrimitive.Root>) {
  const {id, error} = useFormField()

  return (
    <Label
      data-slot="form-label"
      data-error={!!error}
      className={merge('data-[error=true]:text-destructive', className)}
      htmlFor={id}
      {...props}
    />
  )
}

function FormDescription({className, ...props}: ComponentProps<'p'>) {
  const {id} = useFormField()

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

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

  if (!body) {
    return null
  }

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

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