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";
import ProviderSummary from "../components/ProviderSummary";
import OrganizationCard from "../components/OrganizationCard";
import { getRankingBadge } from "../utils/ranking";

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: ["CohereLabs"] },
  { 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, rank }: { provider: ProviderInfo, calendarData: CalendarData, rank: number }) => {
  const providerName = provider.fullName || provider.authors[0];
  const calendarKey = provider.authors[0]; // Use the same key as calendar generation
  const totalCount = calendarData[calendarKey]?.reduce((sum, day) => sum + day.count, 0) || 0;

  const rankingBadge = getRankingBadge(rank);
  
  return (
    <div id={`provider-${calendarKey}`} className={`relative bg-gradient-to-br from-card to-card/95 rounded-2xl border border-border shadow-lg hover:shadow-xl transition-all duration-300 p-6 group ${rank <= 3 ? 'ring-2 ring-yellow-200/50 shadow-yellow-100/20' : ''}`}>
      {/* Ranking Badge */}
      <div className={rankingBadge.className}>
        <span className="flex items-center justify-center">
          {rankingBadge.icon || rank}
        </span>
      </div>
      
      {/* Organization Header */}
      <OrganizationCard
        provider={provider}
        calendarKey={calendarKey}
        providerName={providerName}
        totalCount={totalCount}
      />
      
      {/* Heatmap Section */}
      <div className="flex flex-col items-center bg-muted/30 rounded-xl p-3 group-hover:bg-muted/40 transition-colors duration-300">
        <Heatmap
          data={calendarData[calendarKey]}
          color={provider.color}
          providerName={providerName}
          fullName={provider.fullName ?? providerName}
          avatarUrl={provider.avatarUrl ?? ''}
          authorId={calendarKey}
          showHeader={false}
        />
      </div>
    </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) => {
      const aData = calendarData[a.authors[0]] || [];
      const bData = calendarData[b.authors[0]] || [];
      const aCount = aData.reduce((sum, day) => sum + day.count, 0);
      const bCount = bData.reduce((sum, day) => sum + day.count, 0);
      return bCount - aCount;
    }),
    [providers, calendarData]
  );

  return (
    <div className="w-full p-4 py-16 relative">
      <h1 className="text-3xl lg:text-5xl mt-16 font-bold text-center mb-2 text-foreground">
        Hugging Face Heatmap 🤗
      </h1>
      <div className="text-center text-sm my-8 space-y-4">
        <p className="text-muted-foreground">
          Open models, datasets, and apps from the top AI labs.
        </p>
      </div>

      {/* Provider Summary List */}
      <div className="mb-16 mx-auto">
        <div className="overflow-x-auto scrollbar-hide">
          <div className="flex gap-6 px-4 py-2 min-w-max justify-center">
            {sortedProviders.map((provider, index) => (
              <ProviderSummary
                key={provider.fullName || provider.authors[0]}
                provider={provider}
                calendarData={calendarData}
                rank={index + 1}
              />
            ))}
          </div>
        </div>
      </div>
           
      {isLoading ? (
        <div className="flex items-center justify-center py-20">
          <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-blue-500"></div>
          <p className="ml-4 text-muted-foreground">Loading heatmaps...</p>
        </div>
      ) : (
        <>
          {/* Top 3 Podium */}
          <div className="mb-16">
            <div className="flex flex-col gap-8 max-w-4xl mx-auto">
              {sortedProviders.slice(0, 3).map((provider, index) => (
                <ProviderHeatmap 
                  key={provider.fullName || provider.authors[0]}
                  provider={provider} 
                  calendarData={calendarData}
                  rank={index + 1}
                />
              ))}
            </div>
          </div>

          {/* Rest of the providers */}
          {sortedProviders.length > 3 && (
            <div>
              <h2 className="text-xl font-semibold text-center mb-8 text-muted-foreground">All Organizations</h2>
              <div className="flex flex-col gap-8 max-w-4xl mx-auto">
                {sortedProviders.slice(3).map((provider, index) => (
                  <ProviderHeatmap 
                    key={provider.fullName || provider.authors[0]}
                    provider={provider} 
                    calendarData={calendarData}
                    rank={index + 4}
                  />
                ))}
              </div>
            </div>
          )}
        </>
      )}
      
      {/* CTA Section */}
      <div className="mt-24 mb-16 flex justify-center">
        <div className="bg-gradient-to-br from-card to-card/95 rounded-2xl border border-border shadow-lg hover:shadow-xl transition-all duration-300 p-8 max-w-2xl w-full text-center space-y-6">
          <div className="space-y-4">
            <h2 className="text-2xl lg:text-3xl font-bold text-foreground">
              Want Your Own Heatmap?
            </h2>
            <p className="text-muted-foreground text-lg">
              Search for any Hugging Face organization or user to see their model release activity.
            </p>
          </div>
          <div className="flex justify-center">
            <UserSearchDialog />
          </div>
        </div>
      </div>
    </div>
  );
};

export default OpenSourceHeatmap;