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

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

import React, {BaseSyntheticEvent, 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>,
> = {
  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-fail={!!fieldState.error}
              className={merge('data-[error=true]:text-fail')}
              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-fail 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-fail={!!error}
      className={merge('data-[error=true]:text-fail', 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-fail text-sm', className)}
      {...props}
    >
      {body}
    </p>
  )
}

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