CA/frontend/src/components/CalendarHeatmap.tsx

import React, { useMemo } from 'react';

interface CalendarHeatmapProps {
  data: Array<{ date: string; value: number }>;
  year: number;
  onDayClick?: (date: string) => void;
}

function getColor(value: number): string {
  if (value < 50) return '#10B981';
  if (value < 100) return '#F59E0B';
  if (value < 150) return '#F97316';
  if (value < 200) return '#EF4444';
  return '#7C3AED';
}

function getLabel(value: number): string {
  if (value < 50) return '优';
  if (value < 100) return '良';
  if (value < 150) return '轻度';
  if (value < 200) return '中度';
  return '重度';
}

const MONTH_NAMES = [
  '1月', '2月', '3月', '4月', '5月', '6月',
  '7月', '8月', '9月', '10月', '11月', '12月',
];

const DAY_NAMES = ['一', '二', '三', '四', '五', '六', '日'];

export const CalendarHeatmap = React.memo(function CalendarHeatmap({
  data,
  year,
  onDayClick,
}: CalendarHeatmapProps) {
  const dataMap = useMemo(() => {
    const map = new Map<string, number>();
    for (const d of data) {
      map.set(d.date, d.value);
    }
    return map;
  }, [data]);

  const months = useMemo(() => {
    const result: Array<{
      month: number;
      name: string;
      weeks: Array<Array<{ date: string; day: number; value: number | null }>>;
    }> = [];

    for (let m = 0; m < 12; m++) {
      const daysInMonth = new Date(year, m + 1, 0).getDate();
      const cells: Array<{ date: string; day: number; value: number | null }> = [];

      for (let d = 1; d <= daysInMonth; d++) {
        const dateObj = new Date(year, m, d);
        const dateStr = dateObj.toISOString().slice(0, 10);
        cells.push({
          date: dateStr,
          day: d,
          value: dataMap.get(dateStr) ?? null,
        });
      }

      // Calculate start day of week (1=Monday, 0=Sunday → JS getDay: 0=Sun)
      const firstDay = new Date(year, m, 1).getDay();
      // Convert JS Sunday=0 to Monday=0
      const startOffset = firstDay === 0 ? 6 : firstDay - 1;

      // Pad beginning with empty cells
      const padded: Array<{ date: string; day: number; value: number | null } | null> = [];
      for (let i = 0; i < startOffset; i++) {
        padded.push(null);
      }
      for (const cell of cells) {
        padded.push(cell);
      }

      // Split into weeks of 7
      const weeks: Array<Array<{ date: string; day: number; value: number | null }>> = [];
      for (let i = 0; i < padded.length; i += 7) {
        const week = padded.slice(i, i + 7).filter(Boolean) as Array<{
          date: string;
          day: number;
          value: number | null;
        }>;
        if (week.length > 0) {
          weeks.push(week);
        }
      }

      result.push({ month: m, name: MONTH_NAMES[m], weeks });
    }
    return result;
  }, [year, dataMap]);

  return (
    <div className="space-y-4">
      {/* Legend */}
      <div className="flex items-center gap-1.5 text-[10px] text-gray-500">
        {[
          { color: '#10B981', label: '优 <50' },
          { color: '#F59E0B', label: '良 50-100' },
          { color: '#F97316', label: '轻度 100-150' },
          { color: '#EF4444', label: '中度 150-200' },
          { color: '#7C3AED', label: '重度 ≥200' },
          { color: '#E5E7EB', label: '无数据' },
        ].map((item) => (
          <span key={item.label} className="inline-flex items-center gap-1">
            <span
              className="inline-block w-2.5 h-2.5 rounded-sm"
              style={{ backgroundColor: item.color }}
            />
            {item.label}
          </span>
        ))}
      </div>

      {/* Calendar grid */}
      <div className="grid grid-cols-3 sm:grid-cols-4 md:grid-cols-6 lg:grid-cols-12 gap-4">
        {months.map((month) => (
          <div key={month.month} className="flex flex-col items-center">
            <div className="text-[11px] font-medium text-gray-500 mb-1">
              {month.name}
            </div>
            {/* Day header row */}
            <div className="grid grid-cols-7 gap-px mb-0.5" style={{ width: 98 }}>
              {DAY_NAMES.map((d) => (
                <div
                  key={d}
                  className="text-[8px] text-gray-400 text-center leading-3 w-[14px] h-3"
                >
                  {d}
                </div>
              ))}
            </div>
            {/* Weeks */}
            {month.weeks.map((week, wi) => (
              <div key={wi} className="flex gap-px">
                {Array.from({ length: 7 }).map((_, di) => {
                  // Match by day-of-week index
                  const dayOfWeekMap = [1, 2, 3, 4, 5, 6, 0]; // Mon=1..Sun=0
                  const matchedCell = week.find(
                    (c) => new Date(c.date).getDay() === dayOfWeekMap[di],
                  );

                  if (!matchedCell) {
                    return (
                      <div
                        key={di}
                        className="w-[14px] h-[14px]"
                        aria-hidden
                      />
                    );
                  }

                  const bg =
                    matchedCell.value === null
                      ? '#E5E7EB'
                      : getColor(matchedCell.value);

                  return (
                    <div
                      key={di}
                      title={
                        matchedCell.value !== null
                          ? `${matchedCell.date}: AQI ${matchedCell.value} (${getLabel(matchedCell.value)})`
                          : `${matchedCell.date}: 无数据`
                      }
                      onClick={() => onDayClick?.(matchedCell.date)}
                      className={`w-[14px] h-[14px] rounded-sm transition-transform hover:scale-125 ${
                        onDayClick ? 'cursor-pointer' : ''
                      }`}
                      style={{ backgroundColor: bg }}
                      role={onDayClick ? 'button' : undefined}
                      tabIndex={onDayClick ? 0 : undefined}
                      onKeyDown={
                        onDayClick
                          ? (e: React.KeyboardEvent) => {
                              if (e.key === 'Enter' || e.key === ' ') {
                                e.preventDefault();
                                onDayClick(matchedCell.date);
                              }
                            }
                          : undefined
                      }
                    />
                  );
                })}
              </div>
            ))}
          </div>
        ))}
      </div>
    </div>
  );
});