"use client";

import * as React from "react";
import { cn } from "@/lib/utils";

type RadioGroupContextValue = {
  value: string | undefined;
  setValue: (v: string) => void;
  name: string;
  disabled?: boolean;
  registerItem: (value: string, el: HTMLButtonElement | null) => void;
  unregisterItem: (value: string) => void;
  getItems: () => Array<{ value: string; el: HTMLButtonElement }>;
};

const RadioGroupCtx = React.createContext<RadioGroupContextValue | null>(null);

function useRadioGroupContext() {
  const ctx = React.useContext(RadioGroupCtx);
  if (!ctx)
    throw new Error("RadioGroup primitives must be used inside <RadioGroup>");
  return ctx;
}

export function RadioGroup({
  value: controlledValue,
  defaultValue,
  onValueChange,
  disabled,
  name,
  className,
  children,
  orientation = "vertical",
}: {
  value?: string;
  defaultValue?: string;
  onValueChange?: (v: string) => void;
  disabled?: boolean;
  name?: string;
  className?: string;
  children: React.ReactNode;
  orientation?: "horizontal" | "vertical";
}) {
  const [internal, setInternal] = React.useState<string | undefined>(
    defaultValue,
  );
  const isControlled = controlledValue !== undefined;
  const value = isControlled ? controlledValue : internal;
  const reactId = React.useId();
  const groupName = name ?? `radio-${reactId}`;
  const itemsRef = React.useRef<Map<string, HTMLButtonElement>>(new Map());

  const setValue = React.useCallback(
    (v: string) => {
      if (!isControlled) setInternal(v);
      onValueChange?.(v);
    },
    [isControlled, onValueChange],
  );

  const registerItem = React.useCallback(
    (v: string, el: HTMLButtonElement | null) => {
      if (el) itemsRef.current.set(v, el);
    },
    [],
  );
  const unregisterItem = React.useCallback((v: string) => {
    itemsRef.current.delete(v);
  }, []);
  const getItems = React.useCallback(() => {
    return Array.from(itemsRef.current.entries()).map(([v, el]) => ({
      value: v,
      el,
    }));
  }, []);

  const ctxValue = React.useMemo(
    () => ({
      value,
      setValue,
      name: groupName,
      disabled,
      registerItem,
      unregisterItem,
      getItems,
    }),
    [value, setValue, groupName, disabled, registerItem, unregisterItem, getItems],
  );

  return (
    <RadioGroupCtx.Provider value={ctxValue}>
      <div
        role="radiogroup"
        aria-orientation={orientation}
        className={cn(
          "flex gap-2",
          orientation === "vertical" ? "flex-col" : "flex-row flex-wrap",
          className,
        )}
      >
        {children}
      </div>
    </RadioGroupCtx.Provider>
  );
}

export function RadioGroupItem({
  value,
  id,
  disabled: itemDisabled,
  className,
  children,
}: {
  value: string;
  id?: string;
  disabled?: boolean;
  className?: string;
  children?: React.ReactNode;
}) {
  const ctx = useRadioGroupContext();
  const {
    value: selected,
    setValue,
    disabled: groupDisabled,
    registerItem,
    unregisterItem,
    getItems,
  } = ctx;
  const disabled = itemDisabled || groupDisabled;
  const isSelected = selected === value;
  const reactId = React.useId();
  const inputId = id ?? `radio-item-${reactId}`;
  const buttonRef = React.useRef<HTMLButtonElement | null>(null);

  React.useEffect(() => {
    registerItem(value, buttonRef.current);
    return () => unregisterItem(value);
  }, [value, registerItem, unregisterItem]);

  const handleKeyDown = (e: React.KeyboardEvent<HTMLButtonElement>) => {
    if (disabled) return;
    const items = getItems().filter(
      ({ el }) => !el.hasAttribute("data-disabled"),
    );
    if (items.length === 0) return;
    const currentIdx = items.findIndex((i) => i.value === value);
    let nextIdx = currentIdx;
    if (e.key === "ArrowDown" || e.key === "ArrowRight") {
      e.preventDefault();
      nextIdx = (currentIdx + 1) % items.length;
    } else if (e.key === "ArrowUp" || e.key === "ArrowLeft") {
      e.preventDefault();
      nextIdx = (currentIdx - 1 + items.length) % items.length;
    } else if (e.key === " " || e.key === "Enter") {
      e.preventDefault();
      setValue(value);
      return;
    } else {
      return;
    }
    const next = items[nextIdx];
    if (next) {
      next.el.focus();
      setValue(next.value);
    }
  };

  return (
    <label
      htmlFor={inputId}
      className={cn(
        "flex items-center gap-2",
        disabled
          ? "cursor-not-allowed opacity-50"
          : "cursor-pointer",
      )}
    >
      <button
        ref={buttonRef}
        id={inputId}
        type="button"
        role="radio"
        aria-checked={isSelected}
        disabled={disabled}
        data-disabled={disabled ? "true" : undefined}
        tabIndex={isSelected || (!selected && getFirstFocusableIndex(ctx, value)) ? 0 : -1}
        onClick={() => !disabled && setValue(value)}
        onKeyDown={handleKeyDown}
        className={cn(
          "flex h-4 w-4 shrink-0 items-center justify-center rounded-full border border-input bg-background transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2",
          isSelected && "border-[var(--sag-brand)]",
          className,
        )}
      >
        {isSelected && (
          <span
            aria-hidden="true"
            className="block h-2 w-2 rounded-full bg-[var(--sag-brand)]"
          />
        )}
      </button>
      {children && (
        <span className="text-sm text-foreground select-none">{children}</span>
      )}
    </label>
  );
}

function getFirstFocusableIndex(
  ctx: RadioGroupContextValue,
  value: string,
): boolean {
  // When nothing is selected, only the first item should be in the tab order.
  if (ctx.value !== undefined) return false;
  const items = ctx.getItems();
  return items.length > 0 && items[0].value === value;
}

export function RadioGroupLabel({
  children,
  className,
}: {
  children: React.ReactNode;
  className?: string;
}) {
  return (
    <div
      className={cn(
        "text-sm font-medium text-foreground mb-1.5",
        className,
      )}
    >
      {children}
    </div>
  );
}