Calendar Heatmap

'use client';

import * as React from 'react';

import { cn } from '@/lib/utils';
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/registry/hirael/bases/radix/ui/tooltip';

const CELL_SIZES = { sm: 10, md: 12, lg: 16 } as const;

const DEFAULT_LEVEL_CLASSES = ['bg-muted', 'bg-primary/25', 'bg-primary/50', 'bg-primary/75', 'bg-primary'];

const defaultClassForLevel = (level: number, levels: number) => {
  if (level <= 0 || levels <= 1) return DEFAULT_LEVEL_CLASSES[0];
  const index = Math.max(1, Math.round((level / (levels - 1)) * 4));
  return DEFAULT_LEVEL_CLASSES[Math.min(index, 4)];
};

const toLocalDate = (input: Date | string): Date => {
  if (input instanceof Date) {
    return new Date(input.getFullYear(), input.getMonth(), input.getDate());
  }
  const match = /^(\d{4})-(\d{2})-(\d{2})/.exec(input);
  if (match) {
    return new Date(Number(match[1]), Number(match[2]) - 1, Number(match[3]));
  }
  const parsed = new Date(input);
  return new Date(parsed.getFullYear(), parsed.getMonth(), parsed.getDate());
};

const addDays = (date: Date, days: number): Date => {
  return new Date(date.getFullYear(), date.getMonth(), date.getDate() + days);
};

const addMonthsClamped = (date: Date, months: number): Date => {
  const year = date.getFullYear();
  const month = date.getMonth() + months;
  const lastDay = new Date(year, month + 1, 0).getDate();
  return new Date(year, month, Math.min(date.getDate(), lastDay));
};

const dayKey = (date: Date): number => {
  return date.getFullYear() * 10000 + (date.getMonth() + 1) * 100 + date.getDate();
};

const resolveCellSize = (cellSize: 'sm' | 'md' | 'lg' | number): number => {
  return typeof cellSize === 'number' ? cellSize : CELL_SIZES[cellSize];
};

interface CalendarHeatmapDatum {
  date: Date | string;
  value: number;
}

interface CalendarHeatmapProps extends Omit<React.ComponentProps<'div'>, 'onSelect'> {
  data: CalendarHeatmapDatum[];
  endDate?: Date;
  startDate?: Date;
  months?: number;
  weekStartsOn?: 0 | 1;
  levels?: number;
  thresholds?: number[];
  classForLevel?: (level: number) => string;
  cellSize?: 'sm' | 'md' | 'lg' | number;
  gap?: number;
  locale?: string;
  showMonthLabels?: boolean;
  showWeekdayLabels?: boolean;
  tooltipFormatter?: (date: Date, value: number) => React.ReactNode;
  onSelectDay?: (date: Date, value: number) => void;
}

const CalendarHeatmap = ({
  data,
  endDate,
  startDate,
  months = 12,
  weekStartsOn = 0,
  levels = 5,
  thresholds,
  classForLevel,
  cellSize = 'md',
  gap = 3,
  locale,
  showMonthLabels = true,
  showWeekdayLabels = true,
  tooltipFormatter,
  onSelectDay,
  className,
  ...props
}: CalendarHeatmapProps) => {
  const size = resolveCellSize(cellSize);
  const today = React.useMemo(() => toLocalDate(new Date()), []);

  const { cells, weekCount, monthLabels, weekdayLabels } = React.useMemo(() => {
    const rangeEnd = endDate ? toLocalDate(endDate) : today;
    const rangeStart = startDate ? toLocalDate(startDate) : addDays(addMonthsClamped(rangeEnd, -months), 1);

    const values = new Map<number, number>();
    for (const datum of data) {
      const key = dayKey(toLocalDate(datum.date));
      values.set(key, (values.get(key) ?? 0) + datum.value);
    }

    const leading = (rangeStart.getDay() - weekStartsOn + 7) % 7;
    const gridStart = addDays(rangeStart, -leading);
    const totalDays = Math.round((rangeEnd.getTime() - gridStart.getTime()) / 86400000) + 1;
    const weeks = Math.ceil(totalDays / 7);

    let max = 0;
    for (let i = 0; i < weeks * 7; i++) {
      const date = addDays(gridStart, i);
      if (date < rangeStart || date > rangeEnd) continue;
      max = Math.max(max, values.get(dayKey(date)) ?? 0);
    }
    if (max <= 0) max = 1;

    const cuts =
      thresholds ??
      Array.from({ length: Math.max(levels - 1, 0) }, (_, i) => (max * (i + 1)) / Math.max(levels - 1, 1));

    const levelFor = (value: number) => {
      if (value <= 0) return 0;
      let level = 0;
      for (const cut of cuts) {
        if (value >= cut) level += 1;
      }
      return Math.min(level, levels - 1);
    };

    const dayFormatter = new Intl.DateTimeFormat(locale, { weekday: 'short' });
    const monthFormatter = new Intl.DateTimeFormat(locale, { month: 'short' });

    const grid: {
      date: Date;
      value: number;
      level: number;
      inRange: boolean;
    }[] = [];
    for (let i = 0; i < weeks * 7; i++) {
      const date = addDays(gridStart, i);
      const inRange = date >= rangeStart && date <= rangeEnd;
      const value = inRange ? (values.get(dayKey(date)) ?? 0) : 0;
      grid.push({ date, value, level: inRange ? levelFor(value) : 0, inRange });
    }

    const candidates: { weekIndex: number; label: string }[] = [];
    let lastMonth = -1;
    for (let i = 0; i < grid.length; i++) {
      const cell = grid[i];
      if (!cell.inRange) continue;
      const month = cell.date.getMonth();
      if (month !== lastMonth) {
        lastMonth = month;
        candidates.push({
          weekIndex: Math.floor(i / 7),
          label: monthFormatter.format(cell.date),
        });
      }
    }
    const labels = candidates.filter((candidate, i) => {
      const next = candidates[i + 1];
      return !next || next.weekIndex - candidate.weekIndex >= 3;
    });

    const weekdays = Array.from({ length: 7 }, (_, row) => {
      const day = (weekStartsOn + row) % 7;
      return {
        label: dayFormatter.format(new Date(2024, 0, 7 + day)),
        visible: day === 1 || day === 3 || day === 5,
      };
    });

    return {
      cells: grid,
      weekCount: weeks,
      monthLabels: labels,
      weekdayLabels: weekdays,
    };
  }, [data, endDate, startDate, months, weekStartsOn, levels, thresholds, locale, today]);

  const resolveLevelClass = classForLevel ?? ((level: number) => defaultClassForLevel(level, levels));

  const defaultFormatter = React.useCallback(
    (date: Date, value: number) => {
      const formattedDate = new Intl.DateTimeFormat(locale, {
        month: 'short',
        day: 'numeric',
        year: 'numeric',
      }).format(date);
      const formattedValue = new Intl.NumberFormat(locale).format(value);
      return `${formattedValue} · ${formattedDate}`;
    },
    [locale],
  );

  const formatTooltip = tooltipFormatter ?? defaultFormatter;

  const cellRefs = React.useRef(new Map<number, HTMLButtonElement | null>());
  const [focusedIndex, setFocusedIndex] = React.useState<number | null>(null);

  const defaultTabbableIndex = React.useMemo(() => {
    const todayIndex = cells.findIndex((cell) => cell.inRange && dayKey(cell.date) === dayKey(today));
    if (todayIndex !== -1) return todayIndex;
    for (let i = cells.length - 1; i >= 0; i--) {
      if (cells[i].inRange) return i;
    }
    return -1;
  }, [cells, today]);

  const tabbableIndex = focusedIndex !== null && cells[focusedIndex]?.inRange ? focusedIndex : defaultTabbableIndex;

  const handleCellKeyDown = (event: React.KeyboardEvent<HTMLButtonElement>, index: number) => {
    const rtl = getComputedStyle(event.currentTarget).direction === 'rtl';
    let target: number;
    switch (event.key) {
      case 'ArrowUp':
        target = index - 1;
        break;
      case 'ArrowDown':
        target = index + 1;
        break;
      case 'ArrowLeft':
        target = index + (rtl ? 7 : -7);
        break;
      case 'ArrowRight':
        target = index + (rtl ? -7 : 7);
        break;
      case 'Home':
        target = index - (index % 7);
        while (target < index && !cells[target]?.inRange) target += 1;
        break;
      case 'End':
        target = index - (index % 7) + 6;
        while (target > index && !cells[target]?.inRange) target -= 1;
        break;
      default:
        return;
    }
    event.preventDefault();
    if (target < 0 || target >= cells.length || !cells[target].inRange) return;
    setFocusedIndex(target);
    cellRefs.current.get(target)?.focus();
  };

  return (
    <div data-slot="calendar-heatmap" className={cn('w-fit', className)} {...props}>
      <TooltipProvider delayDuration={0} skipDelayDuration={0}>
        <div className="flex flex-col" style={{ gap }}>
          {showMonthLabels && (
            <div className="flex" style={{ gap }}>
              {showWeekdayLabels && <div aria-hidden className="w-8 shrink-0" />}
              <div
                data-slot="calendar-heatmap-months"
                className="grid"
                style={{
                  gridTemplateColumns: `repeat(${weekCount}, ${size}px)`,
                  columnGap: gap,
                }}
              >
                {monthLabels.map((month) => (
                  <span
                    key={month.weekIndex}
                    data-slot="calendar-heatmap-month-label"
                    className="whitespace-nowrap font-mono text-[10px] uppercase tracking-[0.1em] text-muted-foreground"
                    style={{ gridColumnStart: month.weekIndex + 1, gridRow: 1 }}
                  >
                    {month.label}
                  </span>
                ))}
              </div>
            </div>
          )}
          <div className="flex" style={{ gap }}>
            {showWeekdayLabels && (
              <div
                data-slot="calendar-heatmap-weekdays"
                className="grid w-8 shrink-0"
                style={{
                  gridTemplateRows: `repeat(7, ${size}px)`,
                  rowGap: gap,
                }}
              >
                {weekdayLabels.map((weekday, row) => (
                  <span
                    key={row}
                    data-slot="calendar-heatmap-weekday-label"
                    className={cn(
                      'flex items-center text-[10px] leading-none text-muted-foreground',
                      !weekday.visible && 'invisible',
                    )}
                  >
                    {weekday.label}
                  </span>
                ))}
              </div>
            )}
            <div
              data-slot="calendar-heatmap-grid"
              className="grid grid-flow-col"
              style={{
                gridTemplateRows: `repeat(7, ${size}px)`,
                gridTemplateColumns: `repeat(${weekCount}, ${size}px)`,
                gap,
              }}
            >
              {cells.map((cell, index) => {
                if (!cell.inRange) {
                  return (
                    <div
                      key={cell.date.getTime()}
                      data-slot="calendar-heatmap-cell"
                      aria-hidden
                      className="invisible"
                    />
                  );
                }
                const tooltip = formatTooltip(cell.date, cell.value);
                const label = typeof tooltip === 'string' ? tooltip : defaultFormatter(cell.date, cell.value);
                const cellClassName = cn(
                  'rounded-[2px]',
                  resolveLevelClass(cell.level),
                  onSelectDay &&
                    'cursor-pointer outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-background',
                );
                return (
                  <Tooltip key={cell.date.getTime()}>
                    <TooltipTrigger asChild>
                      {onSelectDay ? (
                        <button
                          type="button"
                          ref={(el) => {
                            cellRefs.current.set(index, el);
                          }}
                          data-slot="calendar-heatmap-cell"
                          data-level={cell.level}
                          aria-label={label}
                          tabIndex={index === tabbableIndex ? 0 : -1}
                          className={cellClassName}
                          onFocus={() => setFocusedIndex(index)}
                          onKeyDown={(event) => handleCellKeyDown(event, index)}
                          onClick={() => onSelectDay(cell.date, cell.value)}
                        />
                      ) : (
                        <div
                          role="img"
                          data-slot="calendar-heatmap-cell"
                          data-level={cell.level}
                          aria-label={label}
                          className={cellClassName}
                        />
                      )}
                    </TooltipTrigger>
                    <TooltipContent>{tooltip}</TooltipContent>
                  </Tooltip>
                );
              })}
            </div>
          </div>
        </div>
      </TooltipProvider>
    </div>
  );
};

interface CalendarHeatmapLegendProps extends React.ComponentProps<'div'> {
  levels?: number;
  classForLevel?: (level: number) => string;
  cellSize?: 'sm' | 'md' | 'lg' | number;
  gap?: number;
  lessLabel?: React.ReactNode;
  moreLabel?: React.ReactNode;
}

const CalendarHeatmapLegend = ({
  levels = 5,
  classForLevel,
  cellSize = 'md',
  gap = 3,
  lessLabel = 'Less',
  moreLabel = 'More',
  className,
  ...props
}: CalendarHeatmapLegendProps) => {
  const size = resolveCellSize(cellSize);
  const resolveLevelClass = classForLevel ?? ((level: number) => defaultClassForLevel(level, levels));

  return (
    <div
      data-slot="calendar-heatmap-legend"
      className={cn('flex w-fit items-center gap-1.5 text-[10px] text-muted-foreground', className)}
      {...props}
    >
      <span data-slot="calendar-heatmap-legend-label">{lessLabel}</span>
      <div className="flex" style={{ gap }}>
        {Array.from({ length: levels }, (_, level) => (
          <span
            key={level}
            data-slot="calendar-heatmap-legend-swatch"
            data-level={level}
            className={cn('rounded-[2px]', resolveLevelClass(level))}
            style={{ width: size, height: size }}
          />
        ))}
      </div>
      <span data-slot="calendar-heatmap-legend-label">{moreLabel}</span>
    </div>
  );
};

export { CalendarHeatmap, CalendarHeatmapLegend };
export type { CalendarHeatmapDatum };