import { useTheme } from "../ThemeProvider";
import { habitProgressColor, themedHabitShade } from "../../lib/habit-palette";
import { useId, useMemo, useLayoutEffect, useRef, useState, type CSSProperties, type ReactNode } from "react";
import { CalendarLegend } from "./CalendarLegend";
import { Field } from "./Field";
import { Button } from "./primitives";
import { CalendarDays } from "lucide-react";
import {
  describeDay,
  calendarTimeline,
  MONTH_VIEWS,
  visibleCalendarDays,
  monthWindow,
  type CalendarDay,
  type MonthView,
} from "./calendar-model";

export function CalendarHeatmap({
  days: allDays,
  unit,
  label,
  color = "#111111",
  compact = false,
  initialMonths = 12,
  selectedDate: controlledDate,
  onSelectDate,
  historyLabel,
  legend,
  selectedDayAction,
  availableRange,
  onRangeChange,
  loading = false,
}: {
  availableRange?: { from: string; to: string };
  onRangeChange?: (range: { from: string; to: string }) => void;
  loading?: boolean;
  days: CalendarDay[];
  unit: string;
  label: string;
  color?: string;
  emptyColor?: string;
  compact?: boolean;
  initialMonths?: MonthView;
  selectedDate?: string;
  onSelectDate?: (date: string) => void;
  historyLabel?: string;
  legend?: ReactNode;
  selectedDayAction?: (date: string) => ReactNode;
}) {
  const { theme } = useTheme();
  const [months, setMonths] = useState<MonthView | "custom">(initialMonths);
  const [customRange, setCustomRange] = useState<{ from: string; to: string } | null>(null);
  const [rangeError, setRangeError] = useState("");
  const latestDate =
    allDays.findLast((day) => day.state !== "future")?.date ??
    allDays.at(-1)?.date;
  // Inspecting a square changes selection, not the window. Only an external
  // date outside the visible preset needs a new anchor to bring it into view.
  const [windowAnchor, setWindowAnchor] = useState({ latestDate, date: latestDate });
  let anchor = windowAnchor.latestDate === latestDate ? windowAnchor.date : latestDate;
  if (months !== "custom" && anchor && controlledDate &&
      allDays.some(day => day.date === controlledDate) &&
      !visibleCalendarDays(allDays, months, anchor).some(day => day.date === controlledDate)) {
    anchor = controlledDate;
  }
  if (windowAnchor.latestDate !== latestDate || windowAnchor.date !== anchor) {
    setWindowAnchor({ latestDate, date: anchor });
  }
  const [localSelectedDate, setLocalSelectedDate] = useState(anchor);
  const selectedDate = controlledDate ?? localSelectedDate;
  function setSelectedDate(date: string) {
    if (onSelectDate) onSelectDate(date);
    else setLocalSelectedDate(date);
  }
  const days = useMemo(() => months === "custom" && customRange
    ? allDays.filter((day) => day.date >= customRange.from && day.date <= customRange.to)
    : anchor ? visibleCalendarDays(allDays, months === "custom" ? 12 : months, anchor) : [], [allDays, months, customRange, anchor]);
  const selected =
    days.find((day) => day.date === selectedDate) ??
    days.find((day) => day.date === anchor) ??
    days[0];
  // Date keys and the shared timeline remain stable when the chart resizes.
  const timeline = useMemo(() => calendarTimeline(days), [days]);
  const buttons = useRef<(HTMLButtonElement | null)[]>([]);
  const inspectorId = useId();
  const instructionsId = useId();
  const rangeErrorId = useId();
  const pendingFocusDate = useRef<string | null>(null);
  const [focusRequest, setFocusRequest] = useState(0);
  const selectedIndex = selected ? days.indexOf(selected) : -1;
  // Controlled editors may reject a selection to preserve an unsaved draft.
  // Only move focus after that date is actually selected, using the new window.
  useLayoutEffect(() => {
    if (pendingFocusDate.current !== null)
      buttons.current[selectedIndex]?.focus();
    pendingFocusDate.current = null;
  }, [focusRequest, selectedIndex, selected?.date]);
  function selectAndFocus(date: string) {
    pendingFocusDate.current = date;
    // A rejected controlled selection still needs a render to restore click
    // focus and consume the request, so a later update cannot reuse it.
    setFocusRequest((request) => request + 1);
    setSelectedDate(date);
  }
  if (!selected) return <p>No calendar dates to display.</p>;
  return (
    <div
      className={`ds-calendar${compact ? " ds-calendar--compact" : ""}`}
      data-months={months}
      aria-busy={loading}
      style={{ "--calendar-weeks": timeline.weeks, "--calendar-color": color } as CSSProperties}
    >
      <div className="ds-calendar-range-toolbar">
        <span className="ds-calendar-range-label" aria-live="polite">
          {new Date(`${days[0]!.date}T12:00:00Z`).toLocaleDateString("en", {
            month: "short",
            day: months === "custom" ? "numeric" : undefined,
            year: "numeric",
            timeZone: "UTC",
          })}{" "}
          {months !== 1 && <>–{" "}
          {new Date(`${days.at(-1)!.date}T12:00:00Z`).toLocaleDateString("en", {
            month: "short",
            day: months === "custom" ? "numeric" : undefined,
            year: "numeric",
            timeZone: "UTC",
          })}
          </>}
        </span>
        <div
          className="ds-month-views"
          role="group"
          aria-label={`${label} time range`}
        >
          {MONTH_VIEWS.map((value) => (
            <button
              key={value}
              type="button"
              aria-label={`${value} ${value === 1 ? "month" : "months"}`}
              aria-pressed={months === value}
              onClick={() => { setMonths(value); if (anchor) onRangeChange?.(monthWindow(anchor, value)); }}
            >
              {value}m
            </button>
          ))}
          <button
            type="button"
            aria-pressed={months === "custom"}
            onClick={() => {
              if (!customRange) setCustomRange({ from: days[0]!.date, to: days.at(-1)!.date });
              setRangeError("");
              setMonths("custom");
            }}
          >
            Custom
          </button>
        </div>
      </div>
      {months === "custom" && customRange && (
        <form
          className="ds-calendar-custom-range"
          aria-label={`${label} custom timeframe`}
          noValidate
          onSubmit={(event) => {
            event.preventDefault();
            const form = event.currentTarget;
            const start = form.elements.namedItem("from") as HTMLInputElement;
            const end = form.elements.namedItem("to") as HTMLInputElement;
            const isAvailable = (input: HTMLInputElement) =>
              input.validity.valid && input.value !== "" &&
              input.value >= (availableRange?.from ?? allDays[0]!.date) && input.value <= (availableRange?.to ?? allDays.at(-1)!.date);
            if (!isAvailable(start) || !isAvailable(end)) {
              setRangeError("Choose both dates within the available history.");
              (!isAvailable(start) ? start : end).focus();
              return;
            }
            if (start.value > end.value) {
              setRangeError("The end date must be on or after the start date.");
              end.focus();
              return;
            }
            if (!availableRange && !allDays.some((day) => day.date >= start.value && day.date <= end.value)) {
              setRangeError("There are no calendar dates in this range.");
              return;
            }
            setCustomRange({ from: start.value, to: end.value });
            onRangeChange?.({ from: start.value, to: end.value });
            setRangeError("");
          }}
        >
          {(["from", "to"] as const).map((name) => (
            <Field key={name} label={name === "from" ? "Start date" : "End date"}>
              {(id) => (
                <input
                  id={id}
                  name={name}
                  type="date"
                  required
                  min={availableRange?.from ?? allDays[0]!.date}
                  max={availableRange?.to ?? allDays.at(-1)!.date}
                  defaultValue={customRange[name]}
                  aria-invalid={rangeError ? true : undefined}
                  aria-describedby={rangeError ? rangeErrorId : undefined}
                  onChange={() => setRangeError("")}
                />
              )}
            </Field>
          ))}
          <Button type="submit" variant="secondary">Apply range</Button>
          {rangeError && <p id={rangeErrorId} className="ds-calendar-range-error" role="alert">{rangeError}</p>}
        </form>
      )}
      {loading && <p role="status" className="ds-supporting-copy">Loading history…</p>}
      <div
        className="ds-calendar-viewport"
        role="group"
        aria-label={label}
        aria-describedby={instructionsId}
      >
        <div className="ds-calendar-timeline">
          <div className="ds-calendar-months" aria-hidden="true">
            {timeline.months.map(({ month, column }, index) => (
              <div className="ds-calendar-month-label" key={month} style={{
                gridColumn: `${column} / ${timeline.months[index + 1]?.column ?? timeline.weeks + 1}`,
              }}>
                {new Date(`${month}-01T12:00:00Z`).toLocaleDateString("en", {
                  month: "short", timeZone: "UTC",
                })}
              </div>
            ))}
          </div>
          <div className="ds-calendar-weekdays" aria-hidden="true">
            <span>Mon</span>
            <span>Wed</span>
            <span>Fri</span>
          </div>
          <div className="ds-calendar-grid">
            {timeline.entries.map(({ day, index, column, row }) => (
              <button
                key={day.date}
                ref={(element) => {
                  buttons.current[index] = element;
                }}
                type="button"
                className={`ds-day ds-day--${day.state}`}
                style={{
                  gridColumn: column,
                  gridRow: row,
                  backgroundColor:
                    day.state === "not-due"
                      ? "transparent"
                      : day.state === "future"
                        ? "var(--ds-empty)"
                        : day.value === 0 ? "var(--ds-empty)" : day.color ? themedHabitShade(day.color, color, theme) : habitProgressColor(color, (Math.max(1, Math.min(8, Math.round(day.value / day.target * 8))) - 1) / 7, theme),
                }}
                aria-label={`${day.date}: ${describeDay(day, unit)}`}
                title={`${day.date}: ${describeDay(day, unit)}`}
                aria-pressed={day.date === selected.date}
                aria-describedby={instructionsId}
                tabIndex={day.date === selected.date ? 0 : -1}
                onClick={() => selectAndFocus(day.date)}
                onKeyDown={(event) => {
                  const offset = {
                    ArrowDown: 1,
                    ArrowUp: -1,
                    ArrowRight: 7,
                    ArrowLeft: -7,
                  }[event.key];
                  const next =
                    event.key === "Home"
                      ? 0
                      : event.key === "End"
                        ? days.length - 1
                        : offset !== undefined
                          ? Math.max(
                              0,
                              Math.min(days.length - 1, index + offset),
                            )
                          : undefined;
                  if (next !== undefined) {
                    event.preventDefault();
                    selectAndFocus(days[next]!.date);
                  }
                }}
              />
            ))}
          </div>
        </div>
      </div>
      {legend ?? <CalendarLegend days={days} color={color} label={label} />}
      <p id={instructionsId} className="ds-visually-hidden">
        Select a day. Up/down: one day. Left/right: one week. Home/end: first/last date.
      </p>
      {!historyLabel && <p className="ds-calendar-demo-label ds-muted">Demo history</p>}
      <div id={inspectorId} className="ds-date-inspector">
        <div className="ds-date-inspector-heading" aria-live="polite">
          <CalendarDays size={22} aria-hidden="true" />
          <div>
            <time dateTime={selected.date}>
              {new Date(`${selected.date}T12:00:00Z`).toLocaleDateString("en", {
                month: "long",
                day: "numeric",
                year: "numeric",
                timeZone: "UTC",
              })}
            </time>
            <span className="ds-date-inspector-status">{describeDay(selected, unit)}</span>
          </div>
        </div>
        {selectedDayAction && <div className="ds-action-row ds-date-inspector-actions">{selectedDayAction(selected.date)}</div>}
      </div>
    </div>
  );
}