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 and item names for each provider and date
  const countMap: Record<string, Record<string, number>> = {};
  const itemsMap: Record<string, Record<string, string[]>> = {};

  modelData.forEach((item) => {
    const dateString = item.createdAt.split("T")[0];
    providers.forEach(({ authors }) => {
      if (authors.some((author) => item.id.startsWith(author))) {
        const providerKey = authors[0];
        
        // Initialize maps if needed
        countMap[providerKey] = countMap[providerKey] || {};
        itemsMap[providerKey] = itemsMap[providerKey] || {};
        
        // Increment count
        countMap[providerKey][dateString] =
          (countMap[providerKey][dateString] || 0) + 1;
          
        // Add item name to the list
        itemsMap[providerKey][dateString] = itemsMap[providerKey][dateString] || [];
        itemsMap[providerKey][dateString].push(item.id);
      }
    });
  });

  // 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 providerKey = authors[0];
      const count = countMap[providerKey]?.[dateString] || 0;
      const items = itemsMap[providerKey]?.[dateString] || [];
      
      data[providerKey].push({ 
        date: dateString, 
        count, 
        level: 0, 
        items 
      });
    });
  }

  // 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;
};