import { createContext, useContext } from 'react';

type FieldLabelContextType = {
  id?: string;
  ref?: React.RefObject<HTMLLabelElement>;
};

const FieldLabelContext = createContext<FieldLabelContextType | undefined>(undefined);
export const FieldLabelContextProvider = FieldLabelContext.Provider;

const InputIdContext = createContext<string | undefined>(undefined);
export const InputIdContextProvider = InputIdContext.Provider;

const InputDescribedByContext = createContext<string | undefined>(undefined);
export const InputDescribedByProvider = InputDescribedByContext.Provider;

const InputInvalidContext = createContext<boolean | undefined>(undefined);
export const InputInvalidProvider = InputInvalidContext.Provider;

interface UseInputAttributesArgs {
  /** Set this to `true` if the underlying element is not directly [labelable as per the HTML specification](https://html.spec.whatwg.org/multipage/forms.html#category-label). */
  nonLabelable?: boolean;
}

export function useInputAttributes({ nonLabelable }: UseInputAttributesArgs = {}) {
  const labelId = useContext(FieldLabelContext)?.id;
  return {
    id: useContext(InputIdContext),
    'aria-labelledby': nonLabelable ? labelId : undefined,
    'aria-describedby': useContext(InputDescribedByContext),
    'aria-invalid': useContext(InputInvalidContext),
  } satisfies React.HTMLAttributes<HTMLElement>;
}

export function useFieldLabelRef() {
  return useContext(FieldLabelContext)?.ref;
}

export interface WithInputAttributesProps {
  inputAttributes: ReturnType<typeof useInputAttributes>;
}

export function withInputAttributes<T extends Partial<WithInputAttributesProps>>(
  Component: React.ComponentType<T>,
  args?: UseInputAttributesArgs,
) {
  function ComponentWithInputAttributes(props: Omit<T, keyof WithInputAttributesProps>) {
    return <Component inputAttributes={useInputAttributes(args)} {...(props as T)} />;
  }

  ComponentWithInputAttributes.displayName = `withInputAttributes(${Component.displayName || Component.name || 'Component'})`;

  return ComponentWithInputAttributes;
}