Date Range Picker

"use client";

import * as React from "react";
import { ChevronLeft, ChevronRight } from "lucide-react";

import { cn } from "@/lib/utils";

export interface DateRange {
  from?: Date;
  to?: Date;
}

export interface DateRangePickerProps {
  value: DateRange;
  onChange: (value: DateRange) => void;
  className?: string;
}

const WEEKDAY_LABELS = ["Su", "Mo", "Tu", "We", "Th", "Fr", "Sa"];

function startOfDay(date: Date) {
  return new Date(date.getFullYear(), date.getMonth(), date.getDate());
}
function isSameDay(a?: Date, b?: Date) {
  return !!a && !!b && a.getFullYear() === b.getFullYear() && a.getMonth() === b.getMonth() && a.getDate() === b.getDate();
}
function isBefore(a: Date, b: Date) {
  return startOfDay(a).getTime() < startOfDay(b).getTime();
}
function addDays(date: Date, days: number) {
  const next = new Date(date);
  next.setDate(next.getDate() + days);
  return next;
}
/** Normalizes to the 1st — used for tracking which month a panel shows. */
function addMonths(date: Date, months: number) {
  return new Date(date.getFullYear(), date.getMonth() + months, 1);
}
/** Keeps the day-of-month (clamped to the target month's length) — used for PageUp/PageDown, which move the focused day, not a panel. */
function addMonthsPreserveDay(date: Date, months: number) {
  const day = date.getDate();
  const targetFirst = addMonths(date, months);
  const daysInTarget = new Date(targetFirst.getFullYear(), targetFirst.getMonth() + 1, 0).getDate();
  return new Date(targetFirst.getFullYear(), targetFirst.getMonth(), Math.min(day, daysInTarget));
}
function isInRange(date: Date, range: DateRange) {
  if (!range.from || !range.to) return false;
  const [start, end] = isBefore(range.from, range.to) ? [range.from, range.to] : [range.to, range.from];
  return !isBefore(date, start) && !isBefore(end, date);
}
function dayId(baseId: string, date: Date) {
  return `${baseId}-day-${date.toISOString().slice(0, 10)}`;
}
function buildMonthGrid(monthDate: Date): (Date | null)[][] {
  const year = monthDate.getFullYear();
  const month = monthDate.getMonth();
  const firstDay = new Date(year, month, 1);
  const daysInMonth = new Date(year, month + 1, 0).getDate();
  const startWeekday = firstDay.getDay();

  const cells: (Date | null)[] = [];
  for (let i = 0; i < startWeekday; i++) cells.push(null);
  for (let d = 1; d <= daysInMonth; d++) cells.push(new Date(year, month, d));
  while (cells.length % 7 !== 0) cells.push(null);

  const weeks: (Date | null)[][] = [];
  for (let i = 0; i < cells.length; i += 7) weeks.push(cells.slice(i, i + 7));
  return weeks;
}

interface MonthGridProps {
  monthDate: Date;
  value: DateRange;
  focusedDate: Date;
  hoveredDate: Date | null;
  today: Date;
  onSelect: (date: Date) => void;
  onFocusDate: (date: Date) => void;
  onHoverDate: (date: Date | null) => void;
  onKeyDown: (event: React.KeyboardEvent, date: Date) => void;
  baseId: string;
}

function MonthGrid({
  monthDate,
  value,
  focusedDate,
  hoveredDate,
  today,
  onSelect,
  onFocusDate,
  onHoverDate,
  onKeyDown,
  baseId,
}: MonthGridProps) {
  const weeks = buildMonthGrid(monthDate);
  const label = monthDate.toLocaleDateString(undefined, { month: "long", year: "numeric" });
  const previewRange: DateRange = value.from && !value.to && hoveredDate ? { from: value.from, to: hoveredDate } : value;

  return (
    <div role="grid" aria-label={label} className="flex flex-col gap-1">
      <p className="px-1 text-sm font-semibold text-foreground">{label}</p>
      <div role="row" className="grid grid-cols-7">
        {WEEKDAY_LABELS.map((day) => (
          <span
            key={day}
            role="columnheader"
            aria-label={day}
            className="flex h-8 items-center justify-center text-xs font-medium text-muted-foreground"
          >
            {day}
          </span>
        ))}
      </div>
      {weeks.map((week, weekIndex) => (
        <div role="row" key={weekIndex} className="grid grid-cols-7">
          {week.map((date, dayIndex) => {
            if (!date) return <span key={dayIndex} aria-hidden="true" />;
            const selectedEdge = isSameDay(date, value.from) || isSameDay(date, value.to);
            const inRange = isInRange(date, previewRange);
            const isFocusTarget = isSameDay(date, focusedDate);
            const isToday = isSameDay(date, today);
            return (
              <div
                key={dayIndex}
                id={dayId(baseId, date)}
                role="gridcell"
                aria-selected={selectedEdge || inRange}
                aria-current={isToday ? "date" : undefined}
                tabIndex={isFocusTarget ? 0 : -1}
                onClick={() => onSelect(date)}
                onFocus={() => onFocusDate(date)}
                onMouseEnter={() => onHoverDate(date)}
                onMouseLeave={() => onHoverDate(null)}
                onKeyDown={(event) => onKeyDown(event, date)}
                className={cn(
                  "flex h-8 w-8 cursor-pointer select-none items-center justify-center rounded-md text-sm outline-none",
                  "focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background",
                  inRange && !selectedEdge && "bg-accent text-accent-foreground",
                  selectedEdge && "bg-primary text-primary-foreground",
                  !selectedEdge && !inRange && "text-foreground hover:bg-accent",
                  isToday && !selectedEdge && "font-semibold",
                )}
              >
                {date.getDate()}
              </div>
            );
          })}
        </div>
      ))}
    </div>
  );
}

/**
 * Two role="grid" calendars sharing one roving-tabindex focus target
 * (focusedDate). Arrow keys move within the currently visible two-month
 * span; crossing past either edge shifts both panels forward/back by one
 * month and refocuses the target day there — this covers real navigation
 * (including PageUp/PageDown month-jumps) without needing an unbounded
 * wrap-around scheme. The live range preview while hovering a potential
 * end date is a mouse-only enhancement; keyboard selection (Enter on the
 * start, then Enter on the end) works identically either way, it just
 * doesn't preview the span before committing.
 */
export function DateRangePicker({ value, onChange, className }: DateRangePickerProps) {
  const today = React.useMemo(() => startOfDay(new Date()), []);
  const [baseMonth, setBaseMonth] = React.useState(() => addMonths(value.from ?? today, 0));
  const [focusedDate, setFocusedDate] = React.useState(() => value.from ?? today);
  const [hoveredDate, setHoveredDate] = React.useState<Date | null>(null);
  const baseId = React.useId();

  const secondMonth = addMonths(baseMonth, 1);
  const rangeStart = startOfDay(baseMonth);
  const rangeEnd = new Date(secondMonth.getFullYear(), secondMonth.getMonth() + 1, 0);

  const selectDate = (date: Date) => {
    if (!value.from || (value.from && value.to)) {
      onChange({ from: date, to: undefined });
    } else {
      const from = value.from;
      onChange(isBefore(date, from) ? { from: date, to: from } : { from, to: date });
    }
  };

  const moveFocus = (next: Date) => {
    setFocusedDate(next);
    if (isBefore(next, rangeStart)) setBaseMonth(addMonths(baseMonth, -1));
    else if (isBefore(rangeEnd, next)) setBaseMonth(addMonths(baseMonth, 1));
    requestAnimationFrame(() => {
      document.getElementById(dayId(baseId, next))?.focus();
    });
  };

  const handleKeyDown = (event: React.KeyboardEvent, date: Date) => {
    let next: Date | null = null;
    if (event.key === "ArrowRight") next = addDays(date, 1);
    else if (event.key === "ArrowLeft") next = addDays(date, -1);
    else if (event.key === "ArrowDown") next = addDays(date, 7);
    else if (event.key === "ArrowUp") next = addDays(date, -7);
    else if (event.key === "Home") next = addDays(date, -date.getDay());
    else if (event.key === "End") next = addDays(date, 6 - date.getDay());
    else if (event.key === "PageDown") next = addMonthsPreserveDay(date, 1);
    else if (event.key === "PageUp") next = addMonthsPreserveDay(date, -1);
    else if (event.key === "Enter" || event.key === " ") {
      event.preventDefault();
      selectDate(date);
      return;
    }
    if (!next) return;
    event.preventDefault();
    moveFocus(next);
  };

  return (
    <div role="group" aria-label="Date range" className={cn("inline-flex flex-col gap-4 rounded-lg border border-border bg-card p-4", className)}>
      <div className="flex items-start gap-4 sm:gap-8">
        <button
          type="button"
          aria-label="Previous month"
          onClick={() => setBaseMonth(addMonths(baseMonth, -1))}
          className="mt-6 flex size-7 shrink-0 items-center justify-center rounded-md text-foreground outline-none hover:bg-accent focus-visible:ring-2 focus-visible:ring-ring"
        >
          <ChevronLeft className="size-4" aria-hidden="true" />
        </button>
        <div className="flex flex-1 flex-col gap-4 sm:flex-row sm:gap-8">
          <MonthGrid
            monthDate={baseMonth}
            value={value}
            focusedDate={focusedDate}
            hoveredDate={hoveredDate}
            today={today}
            onSelect={selectDate}
            onFocusDate={setFocusedDate}
            onHoverDate={setHoveredDate}
            onKeyDown={handleKeyDown}
            baseId={baseId}
          />
          <MonthGrid
            monthDate={secondMonth}
            value={value}
            focusedDate={focusedDate}
            hoveredDate={hoveredDate}
            today={today}
            onSelect={selectDate}
            onFocusDate={setFocusedDate}
            onHoverDate={setHoveredDate}
            onKeyDown={handleKeyDown}
            baseId={baseId}
          />
        </div>
        <button
          type="button"
          aria-label="Next month"
          onClick={() => setBaseMonth(addMonths(baseMonth, 1))}
          className="mt-6 flex size-7 shrink-0 items-center justify-center rounded-md text-foreground outline-none hover:bg-accent focus-visible:ring-2 focus-visible:ring-ring"
        >
          <ChevronRight className="size-4" aria-hidden="true" />
        </button>
      </div>
    </div>
  );
}