src/pages/index.tsx

import React, { useState, useEffect, useMemo } from "react";
import { generateCalendarData } from "../utils/calendar";
import {
  OpenSourceHeatmapProps,
  ProviderInfo,
  ModelData,
  CalendarData,
} from "../types/heatmap";
import Heatmap from "../components/Heatmap";
import { fetchAllProvidersData, fetchAllAuthorsData } from "../utils/authors";
import UserSearchDialog from "../components/UserSearchDialog";

const PROVIDERS: ProviderInfo[] = [
  { color: "#ff7000", authors: ["mistralai"] },
  { color: "#1877F2", authors: ["meta-llama", "facebook", ] },
  { color: "#10A37F", authors: ["openai"] },
  { color: "#cc785c", authors: ["Anthropic"] },
  { color: "#DB4437", authors: ["google"] },
  { color: "#5E35B1", authors: ["allenai"] },
  { color: "#0088cc", authors: ["apple"] },
  { color: "#FEB800", authors: ["microsoft"] },
  { color: "#76B900", authors: ["nvidia"] },
  { color: "#0088cc", authors: ["deepseek-ai"] },
  { color: "#0088cc", authors: ["Qwen"] },
  { color: "#4C6EE6", authors: ["CohereForAI"] },
  { color: "#4C6EE6", authors: ["ibm-granite"] },
  { color: "#A020F0", authors: ["stabilityai"] },
];

export async function getStaticProps() {
  try {
    const allAuthors = PROVIDERS.flatMap(({ authors }) => authors);
    const uniqueAuthors = Array.from(new Set(allAuthors));

    const flatData: ModelData[] = await fetchAllAuthorsData(uniqueAuthors);
    const updatedProviders = await fetchAllProvidersData(PROVIDERS);

    const calendarData = generateCalendarData(flatData, updatedProviders);

    return {
      props: {
        calendarData,
        providers: updatedProviders,
      },
      revalidate: 3600, // regenerate every hour
    };
  } catch (error) {
    console.error("Error fetching data:", error);
    return {
      props: {
        calendarData: {},
        providers: PROVIDERS,
      },
      revalidate: 60, // retry after 1 minute if there was an error
    };
  }
}

const ProviderHeatmap = React.memo(({ provider, calendarData }: { provider: ProviderInfo, calendarData: CalendarData }) => {
  const providerName = provider.fullName || provider.authors[0];
  return (
    <div key={providerName} className="flex flex-col items-center">
      <Heatmap
        data={calendarData[providerName]}
        color={provider.color}
        providerName={providerName}
        fullName={provider.fullName ?? providerName}
        avatarUrl={provider.avatarUrl ?? ''}
      />
    </div>
  );
});

const OpenSourceHeatmap: React.FC<OpenSourceHeatmapProps> = ({
  calendarData,
  providers,
}) => {
  const [isLoading, setIsLoading] = useState(true);

  useEffect(() => {
    if (calendarData && Object.keys(calendarData).length > 0) {
      setIsLoading(false);
    }
  }, [calendarData]);

  const sortedProviders = useMemo(() => 
    providers.sort((a, b) =>
      calendarData[b.fullName || b.authors[0]].reduce(
        (sum, day) => sum + day.count,
        0
      ) -
      calendarData[a.fullName || a.authors[0]].reduce(
        (sum, day) => sum + day.count,
        0
      )
    ),
    [providers, calendarData]
  );

  return (
    <div className="w-full max-w-screen-lg mx-auto p-4 py-16">
      <h1 className="text-3xl lg:text-5xl mt-16 font-bold text-center mb-2">
        Hugging Face Heatmap 🤗
      </h1>
      <div className="text-center text-sm my-8 space-y-4">
        <p>
          Models, Datasets, and Spaces from the top AI labs.
        </p>
        <div className="flex justify-center space-x-4">
          <UserSearchDialog />
        </div>
      </div>
      
      <div className="h-px max-w-lg mx-auto bg-gray-200 my-16" />
     
      {isLoading ? (
        <p className="text-center">Loading...</p>
      ) : (
        <div className="space-y-16">
          {sortedProviders.map((provider) => (
            <ProviderHeatmap 
              key={provider.fullName || provider.authors[0]} 
              provider={provider} 
              calendarData={calendarData} 
            />
          ))}
        </div>
      )}
    </div>
  );
};

export default React.memo(OpenSourceHeatmap);