road-monitoring-ui/src/components/ui/calendar-grid-picker.tsx

'use client';

import * as React from 'react';
import { ChevronLeftIcon, ChevronRightIcon } from 'lucide-react';
import { useDayPicker, type MonthCaptionProps } from 'react-day-picker';

import { cn } from '@/lib/utils';
import { Button } from '@/components/ui/button';
import { Calendar, CalendarDayButton } from '@/components/ui/calendar';

type CalendarView = 'days' | 'months' | 'years';

const MONTH_LABELS = [
  'Jan',
  'Feb',
  'Mar',
  'Apr',
  'May',
  'Jun',
  'Jul',
  'Aug',
  'Sep',
  'Oct',
  'Nov',
  'Dec',
] as const;

const PICKER_WIDTH = 'w-[280px]';
const PICKER_SHELL_CLASS =
  'group/calendar p-3 text-popover-foreground [--cell-size:--spacing(8)]';
const PICKER_HEADER_CLASS =
  'grid min-h-8 grid-cols-[2rem_1fr_2rem] items-center gap-1 border-b border-border pb-2';
const PICKER_NAV_BUTTON_CLASS = 'size-8 shrink-0 p-0';
const GRID_BUTTON_CLASS =
  'h-9 w-full rounded-lg font-medium text-foreground hover:bg-accent/60 hover:text-accent-foreground';
const GRID_BUTTON_ACTIVE_CLASS =
  'bg-muted text-foreground hover:bg-muted hover:text-foreground';
const GRID_BUTTON_DISABLED_CLASS =
  'text-muted-foreground opacity-50 hover:bg-transparent hover:text-muted-foreground';

function getMonthIndex(date: Date) {
  return date.getFullYear() * 12 + date.getMonth();
}

function isYearOutsideRange(year: number, startMonth?: Date, endMonth?: Date) {
  if (startMonth && year < startMonth.getFullYear()) return true;
  if (endMonth && year > endMonth.getFullYear()) return true;
  return false;
}

function isMonthOutsideRange(
  year: number,
  month: number,
  startMonth?: Date,
  endMonth?: Date,
) {
  const monthIndex = year * 12 + month;

  if (startMonth && monthIndex < getMonthIndex(startMonth)) return true;
  if (endMonth && monthIndex > getMonthIndex(endMonth)) return true;

  return false;
}

function clampMonthToRange(
  year: number,
  month: number,
  startMonth?: Date,
  endMonth?: Date,
) {
  let nextMonth = month;

  if (startMonth && year === startMonth.getFullYear()) {
    nextMonth = Math.max(nextMonth, startMonth.getMonth());
  }

  if (endMonth && year === endMonth.getFullYear()) {
    nextMonth = Math.min(nextMonth, endMonth.getMonth());
  }

  return Math.min(11, Math.max(0, nextMonth));
}

type CalendarGridPickerContextValue = {
  view: CalendarView;
  setView: (view: CalendarView) => void;
  pickerYear: number;
  setPickerYear: React.Dispatch<React.SetStateAction<number>>;
  decadeStart: number;
  setDecadeStart: React.Dispatch<React.SetStateAction<number>>;
};

const CalendarGridPickerContext =
  React.createContext<CalendarGridPickerContextValue | null>(null);

function useCalendarGridPicker() {
  const context = React.useContext(CalendarGridPickerContext);
  if (!context) {
    throw new Error(
      'useCalendarGridPicker must be used within CalendarGridPicker',
    );
  }
  return context;
}

function PickerNav({
  label,
  onPrevious,
  onNext,
  onLabelClick,
  disablePrevious = false,
  disableNext = false,
}: {
  label: string;
  onPrevious: () => void;
  onNext: () => void;
  onLabelClick?: () => void;
  disablePrevious?: boolean;
  disableNext?: boolean;
}) {
  return (
    <div className={PICKER_HEADER_CLASS}>
      <Button
        type="button"
        variant="ghost"
        size="icon-sm"
        className={cn(PICKER_NAV_BUTTON_CLASS, 'justify-self-start')}
        onClick={onPrevious}
        disabled={disablePrevious}
        aria-label="Previous"
      >
        <ChevronLeftIcon className="size-4" />
      </Button>
      {onLabelClick ? (
        <button
          type="button"
          className="rounded-lg px-2 text-center text-sm font-semibold text-foreground hover:text-accent-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
          onClick={onLabelClick}
        >
          {label}
        </button>
      ) : (
        <span className="px-2 text-center text-sm font-semibold text-foreground">
          {label}
        </span>
      )}
      <Button
        type="button"
        variant="ghost"
        size="icon-sm"
        className={cn(PICKER_NAV_BUTTON_CLASS, 'justify-self-end')}
        onClick={onNext}
        disabled={disableNext}
        aria-label="Next"
      >
        <ChevronRightIcon className="size-4" />
      </Button>
    </div>
  );
}

function MonthGrid({
  year,
  activeMonth,
  onSelect,
  startMonth,
  endMonth,
}: {
  year: number;
  activeMonth: number;
  onSelect: (month: number) => void;
  startMonth?: Date;
  endMonth?: Date;
}) {
  return (
    <div className="grid grid-cols-3 gap-x-2 gap-y-1.5 pt-2">
      {MONTH_LABELS.map((label, month) => {
        const isDisabled = isMonthOutsideRange(
          year,
          month,
          startMonth,
          endMonth,
        );

        return (
          <Button
            key={label}
            type="button"
            variant="ghost"
            disabled={isDisabled}
            className={cn(
              GRID_BUTTON_CLASS,
              isDisabled && GRID_BUTTON_DISABLED_CLASS,
              activeMonth === month && !isDisabled && GRID_BUTTON_ACTIVE_CLASS,
            )}
            onClick={() => onSelect(month)}
          >
            {label}
          </Button>
        );
      })}
    </div>
  );
}

function YearGrid({
  decadeStart,
  activeYear,
  onSelect,
  startMonth,
  endMonth,
}: {
  decadeStart: number;
  activeYear: number;
  onSelect: (year: number) => void;
  startMonth?: Date;
  endMonth?: Date;
}) {
  const years = React.useMemo(
    () => Array.from({ length: 10 }, (_, index) => decadeStart + index),
    [decadeStart],
  );

  return (
    <div className="grid grid-cols-2 gap-x-3 gap-y-1.5 pt-2">
      {years.map((year) => {
        const isDisabled = isYearOutsideRange(year, startMonth, endMonth);

        return (
          <Button
            key={year}
            type="button"
            variant="ghost"
            disabled={isDisabled}
            className={cn(
              GRID_BUTTON_CLASS,
              isDisabled && GRID_BUTTON_DISABLED_CLASS,
              activeYear === year && !isDisabled && GRID_BUTTON_ACTIVE_CLASS,
            )}
            onClick={() => onSelect(year)}
          >
            {year}
          </Button>
        );
      })}
    </div>
  );
}

function CalendarMonthCaption({
  calendarMonth,
  className,
  children: _children,
  displayIndex: _displayIndex,
  ...props
}: MonthCaptionProps) {
  const { goToMonth, previousMonth, nextMonth } = useDayPicker();
  const { setView, setPickerYear, setDecadeStart } = useCalendarGridPicker();
  const date = calendarMonth.date;
  const monthName = date.toLocaleString('default', { month: 'long' });
  const year = date.getFullYear();

  return (
    <div className={cn(PICKER_HEADER_CLASS, className)} {...props}>
      <Button
        type="button"
        variant="ghost"
        size="icon-sm"
        className={cn(PICKER_NAV_BUTTON_CLASS, 'justify-self-start')}
        disabled={!previousMonth}
        aria-label="Previous month"
        onClick={() => previousMonth && goToMonth(previousMonth)}
      >
        <ChevronLeftIcon className="size-4" />
      </Button>

      <div className="flex items-center justify-center gap-1 text-sm font-semibold">
        <button
          type="button"
          className="rounded-lg px-1 text-center text-sm font-semibold text-foreground hover:text-accent-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
          onClick={() => {
            setPickerYear(year);
            setView('months');
          }}
        >
          {monthName}
        </button>
        <button
          type="button"
          className="rounded-lg px-1 text-center text-sm font-semibold text-foreground hover:text-accent-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
          onClick={() => {
            setDecadeStart(Math.floor(year / 10) * 10);
            setPickerYear(year);
            setView('years');
          }}
        >
          {year}
        </button>
      </div>

      <Button
        type="button"
        variant="ghost"
        size="icon-sm"
        className={cn(PICKER_NAV_BUTTON_CLASS, 'justify-self-end')}
        disabled={!nextMonth}
        aria-label="Next month"
        onClick={() => nextMonth && goToMonth(nextMonth)}
      >
        <ChevronRightIcon className="size-4" />
      </Button>
    </div>
  );
}

function CalendarGridPickerPanel({
  month,
  onMonthChange,
  startMonth,
  endMonth,
}: {
  month: Date;
  onMonthChange: (month: Date) => void;
  startMonth?: Date;
  endMonth?: Date;
}) {
  const {
    view,
    setView,
    pickerYear,
    setPickerYear,
    decadeStart,
    setDecadeStart,
  } = useCalendarGridPicker();

  const minYear = startMonth?.getFullYear();
  const maxYear = endMonth?.getFullYear();
  const activeMonth = clampMonthToRange(
    pickerYear,
    month.getMonth(),
    startMonth,
    endMonth,
  );

  if (view === 'months') {
    return (
      <div className="w-full">
        <PickerNav
          label={String(pickerYear)}
          onPrevious={() => setPickerYear((year) => year - 1)}
          onNext={() => setPickerYear((year) => year + 1)}
          onLabelClick={() => {
            setDecadeStart(Math.floor(pickerYear / 10) * 10);
            setView('years');
          }}
          disablePrevious={minYear !== undefined && pickerYear <= minYear}
          disableNext={maxYear !== undefined && pickerYear >= maxYear}
        />
        <MonthGrid
          year={pickerYear}
          activeMonth={activeMonth}
          startMonth={startMonth}
          endMonth={endMonth}
          onSelect={(nextMonth) => {
            onMonthChange(new Date(pickerYear, nextMonth, 1));
            setView('days');
          }}
        />
      </div>
    );
  }

  if (view === 'years') {
    return (
      <div className="w-full">
        <PickerNav
          label={`${decadeStart} - ${decadeStart + 9}`}
          onPrevious={() => setDecadeStart((start) => start - 10)}
          onNext={() => setDecadeStart((start) => start + 10)}
          disablePrevious={minYear !== undefined && decadeStart <= minYear}
          disableNext={maxYear !== undefined && decadeStart + 10 > maxYear}
        />
        <YearGrid
          decadeStart={decadeStart}
          activeYear={pickerYear}
          startMonth={startMonth}
          endMonth={endMonth}
          onSelect={(year) => {
            const nextMonth = clampMonthToRange(
              year,
              activeMonth,
              startMonth,
              endMonth,
            );

            setPickerYear(year);
            onMonthChange(new Date(year, nextMonth, 1));
            setView('months');
          }}
        />
      </div>
    );
  }

  return null;
}

export type CalendarGridPickerProps = Omit<
  Extract<React.ComponentProps<typeof Calendar>, { mode: 'single' }>,
  'captionLayout' | 'navLayout' | 'hideNavigation'
> & {
  onViewChange?: (view: CalendarView) => void;
};

function CalendarGridPicker({
  className,
  classNames,
  components,
  formatters,
  month,
  defaultMonth,
  onMonthChange,
  onViewChange,
  startMonth,
  endMonth,
  ...props
}: CalendarGridPickerProps) {
  const initialMonth = month ?? defaultMonth ?? new Date();
  const [view, setView] = React.useState<CalendarView>('days');
  const [pickerYear, setPickerYear] = React.useState(
    initialMonth.getFullYear(),
  );
  const [decadeStart, setDecadeStart] = React.useState(
    Math.floor(initialMonth.getFullYear() / 10) * 10,
  );
  const [internalMonth, setInternalMonth] = React.useState(initialMonth);

  const currentMonth = month ?? internalMonth;

  React.useEffect(() => {
    if (month) {
      setPickerYear(month.getFullYear());
      setDecadeStart(Math.floor(month.getFullYear() / 10) * 10);
    }
  }, [month]);

  const handleViewChange = React.useCallback(
    (nextView: CalendarView) => {
      setView(nextView);
      onViewChange?.(nextView);
    },
    [onViewChange],
  );

  const handleMonthChange = React.useCallback(
    (nextMonth: Date) => {
      if (!month) setInternalMonth(nextMonth);
      setPickerYear(nextMonth.getFullYear());
      setDecadeStart(Math.floor(nextMonth.getFullYear() / 10) * 10);
      onMonthChange?.(nextMonth);
    },
    [month, onMonthChange],
  );

  const contextValue = React.useMemo(
    () => ({
      view,
      setView: handleViewChange,
      pickerYear,
      setPickerYear,
      decadeStart,
      setDecadeStart,
    }),
    [view, handleViewChange, pickerYear, decadeStart],
  );

  if (view !== 'days') {
    return (
      <CalendarGridPickerContext.Provider value={contextValue}>
        <div className={cn(PICKER_SHELL_CLASS, PICKER_WIDTH, className)}>
          <CalendarGridPickerPanel
            month={currentMonth}
            onMonthChange={handleMonthChange}
            startMonth={startMonth}
            endMonth={endMonth}
          />
        </div>
      </CalendarGridPickerContext.Provider>
    );
  }

  return (
    <CalendarGridPickerContext.Provider value={contextValue}>
      <Calendar
        hideNavigation
        className={cn(PICKER_SHELL_CLASS, PICKER_WIDTH, className)}
        month={currentMonth}
        onMonthChange={handleMonthChange}
        startMonth={startMonth}
        endMonth={endMonth}
        formatters={{
          formatWeekdayName: (date) =>
            date.toLocaleString('en-US', { weekday: 'short' }).slice(0, 2),
          ...formatters,
        }}
        classNames={{
          month: 'flex w-full flex-col gap-2',
          month_caption: 'p-0',
          weekdays: 'flex w-full',
          weekday:
            'flex-1 text-center text-[0.8rem] font-medium text-muted-foreground select-none',
          week: 'mt-1 flex w-full',
          ...classNames,
        }}
        components={{
          MonthCaption: CalendarMonthCaption,
          DayButton: (dayButtonProps) => (
            <CalendarDayButton
              {...dayButtonProps}
              className={cn(
                dayButtonProps.className,
                'data-[selected-single=true]:rounded-full data-[selected-single=true]:bg-muted data-[selected-single=true]:text-foreground',
              )}
            />
          ),
          ...components,
        }}
        {...(props as Extract<
          React.ComponentProps<typeof Calendar>,
          { mode: 'single' }
        >)}
      />
    </CalendarGridPickerContext.Provider>
  );
}

export { CalendarGridPicker, type CalendarView };