import { composeEventHandlers } from "@radix-ui/primitive"
import { composeRefs } from "@radix-ui/react-compose-refs"
import { Primitive } from "@radix-ui/react-primitive"
import { Slot } from "@radix-ui/react-slot"
import { cva } from "class-variance-authority"
import clsx from "clsx"
import {
  type ComponentProps,
  createContext,
  isValidElement,
  type RefObject,
  useContext,
  useMemo,
  useRef,
  useState,
} from "react"

export type InputBaseContextProps = Pick<InputBaseProps, "autoFocus" | "disabled" | "size"> & {
  controlRef: RefObject<HTMLElement | null>
  onFocusedChange: (focused: boolean) => void
}

const InputContext = createContext<InputBaseContextProps>({
  autoFocus: false,
  controlRef: { current: null },
  disabled: false,
  size: "md",
  onFocusedChange: () => {},
})

const useInputContext = () => useContext(InputContext)

export type InputSize = "sm" | "md" | "lg"

export interface InputBaseProps extends ComponentProps<typeof Primitive.div> {
  autoFocus?: boolean
  disabled?: boolean
  error?: boolean
  size?: InputSize
}

const inputVariants = cva(
  "flex cursor-text items-center gap-2 rounded-md border border-input bg-transparent text-sm ring-offset-background",
  {
    variants: {
      size: {
        sm: "px-2 py-1.5",
        md: "px-3 py-2",
        lg: "px-4 py-3",
      },
      disabled: {
        true: "cursor-not-allowed opacity-50",
      },
      focused: {
        true: "outline-none ring-2 ring-primary",
      },
      error: {
        true: "ring-2 ring-red-500",
      },
    },
    defaultVariants: {
      size: "md",
    },
  },
)

const inputIconVariants = cva("flex items-center text-muted-foreground", {
  variants: {
    size: {
      sm: "[&_svg]:size-3",
      md: "[&_svg]:size-4",
      lg: "[&_svg]:size-4.5",
    },
  },
  defaultVariants: {
    size: "md",
  },
})

const adornmentButtonVariants = cva(
  "inline-flex items-center justify-center rounded-md transition-colors disabled:pointer-events-none disabled:opacity-50 hover:bg-accent hover:text-accent-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring shrink-0",
  {
    variants: {
      size: {
        sm: "size-3 text-xs",
        md: "size-4 text-sm",
        lg: "size-5 text-base",
      },
    },
    defaultVariants: {
      size: "md",
    },
  },
)

export const InputRoot = ({
  autoFocus,
  disabled,
  className,
  onClick,
  error,
  ref,
  size = "md",
  ...props
}: InputBaseProps) => {
  const [focused, setFocused] = useState(false)

  const controlRef = useRef<HTMLElement | null>(null)

  const contextValue = useMemo(
    () => ({
      autoFocus,
      controlRef,
      disabled,
      size,
      onFocusedChange: setFocused,
    }),
    [autoFocus, disabled, size],
  )

  return (
    <InputContext.Provider value={contextValue}>
      <Primitive.div
        ref={ref}
        onClick={composeEventHandlers(onClick, event => {
          // Based on MUI's <InputBase /> implementation.
          // https://github.com/mui/material-ui/blob/master/packages/mui-material/src/InputBase/InputBase.js#L458~L460
          if (controlRef.current && event.currentTarget === event.target) {
            controlRef.current.focus()
          }
        })}
        className={clsx(inputVariants({ size, disabled, focused, error }), className)}
        {...props}
        aria-invalid={Boolean(error)}
      />
    </InputContext.Provider>
  )
}
InputRoot.displayName = "InputRoot"

export const InputFlexWrapper = ({ className, ...props }: ComponentProps<typeof Primitive.div>) => (
  <Primitive.div className={clsx("flex flex-1 flex-wrap", className)} {...props} />
)
InputFlexWrapper.displayName = "InputFlexWrapper"

export const InputControl = ({ onFocus, onBlur, ref, ...props }: ComponentProps<typeof Slot>) => {
  const { controlRef, autoFocus, disabled, onFocusedChange } = useInputContext()

  return (
    <Slot
      ref={composeRefs(controlRef, ref)}
      autoFocus={autoFocus}
      onFocus={composeEventHandlers(onFocus, () => onFocusedChange(true))}
      onBlur={composeEventHandlers(onBlur, () => onFocusedChange(false))}
      {...{ disabled }}
      {...props}
    />
  )
}
InputControl.displayName = "InputControl"

export interface InputBaseAdornmentProps extends ComponentProps<"div"> {
  asChild?: boolean
  disablePointerEvents?: boolean
}

export const InputAdornment = ({
  className,
  disablePointerEvents,
  asChild,
  children,
  ...props
}: InputBaseAdornmentProps) => {
  const { size = "md" } = useInputContext()

  // Extract nested ternary into separate variable
  const childrenComp = typeof children === "string" ? "p" : "div"
  const Comp = asChild ? Slot : childrenComp

  const isAction = isValidElement(children) && children.type === InputAdornmentButton

  return (
    <Comp
      className={clsx(
        inputIconVariants({ size }),
        (!isAction || disablePointerEvents) && "pointer-events-none",
        className,
      )}
      {...props}
    >
      {children}
    </Comp>
  )
}
InputAdornment.displayName = "InputAdornment"

export const InputAdornmentButton = ({
  type = "button",
  disabled: disabledProp,
  className,
  ...props
}: ComponentProps<typeof Primitive.button>) => {
  const { disabled, size: inputSize = "md" } = useInputContext()

  return (
    <Primitive.button
      type={type}
      disabled={disabled ?? disabledProp}
      className={clsx(adornmentButtonVariants({ size: inputSize }), className)}
      {...props}
    />
  )
}

InputAdornmentButton.displayName = "InputAdornmentButton"

// Exclude HTML input 'size' attribute to avoid conflicts with our custom size prop
export interface InputProps extends Omit<ComponentProps<typeof Primitive.input>, "size"> {}

export const Input = ({ className, ...props }: InputProps) => (
  <Primitive.input
    className={clsx(
      "w-full flex-1 bg-transparent file:border-0 file:bg-transparent file:text-sm file:font-medium file:text-foreground placeholder:text-muted-foreground focus:outline-none disabled:pointer-events-none",
      className,
    )}
    {...props}
  />
)
Input.displayName = "Input"