import {
  ModelData,
  ProviderInfo,
  CalendarData,
  Activity,
} from "../types/heatmap";

export const generateCalendarData = (
  modelData: ModelData[],
  providers: ProviderInfo[]
): CalendarData => {
  const data: Record<string, Activity[]> = Object.fromEntries(
    providers.map((provider) => [provider.authors[0], []]),
  );

  const today = new Date();

  // Calculate the first date to display GitHub-style:
  // 1. Go back 52 weeks (364 days) from today so we cover the last year.
  // 2. Shift further back to the previous Sunday (if necessary) so that
  //    the first column of the heatmap always begins on a Sunday.
  const startDate = new Date(today);
  startDate.setDate(startDate.getDate() - 364);
  
  const dayOfWeek = startDate.getDay();
  startDate.setDate(startDate.getDate() - dayOfWeek);

  // create a map to store counts for each provider and date
  const countMap: Record<string, Record<string, number>> = {};

  modelData.forEach((item) => {
    const dateString = item.createdAt.split("T")[0];
    providers.forEach(({ authors }) => {
      if (authors.some((author) => item.id.startsWith(author))) {
        countMap[authors[0]] = countMap[authors[0]] || {};
        countMap[authors[0]][dateString] =
          (countMap[authors[0]][dateString] || 0) + 1;
      }
    });
  });

  // fill in with 0s for days with no activity
  for (let d = new Date(startDate); d <= today; d.setDate(d.getDate() + 1)) {
    const dateString = d.toISOString().split("T")[0];

    providers.forEach(({ authors }) => {
      const count = countMap[authors[0]]?.[dateString] || 0;
      data[authors[0]].push({ date: dateString, count, level: 0 });
    });
  }

  // calculate average counts for each provider
  const avgCounts = Object.fromEntries(
    Object.entries(data).map(([provider, days]) => [
      provider,
      days.reduce((sum, day) => sum + day.count, 0) / days.length || 0,
    ]),
  );

  const calculateLevel = (count: number, avgCount: number): number => {
    if (count === 0) {
      return 0;
    }
    if (count <= avgCount * 0.5) {
      return 1;
    }
    if (count <= avgCount) {
      return 2;
    }
    if (count <= avgCount * 1.5) {
      return 3;
    }
    return 4;
  };

  // assign levels based on count relative to average
  Object.entries(data).forEach(([provider, days]) => {
    const avgCount = avgCounts[provider];
    for (const day of days) {
      day.level = calculateLevel(day.count, avgCount);
    }
  });

  return data;
};