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: "#E91E63", authors: ["Baidu","PaddlePaddle"] },
  { color: "#00BCD4", authors: ["tencent"] },
  { color: "#FFD700", authors: ["BAAI"] },
  { color: "#DC143C", authors: ["OpenGVLab","InternLM"] },
  { color: "#FF4500", authors: ["Skywork"] },
  { color: "#556B2F", authors: ["MiniMaxAI"] },
  { color: "#7FFFD4", authors: ["stepfun-ai"] },
  { color: "#FF6347", authors: ["ByteDance","Bytedance Seed","bytedance-research"]},
  { color: "#9C27B0", authors: ["openbmb"] },
  { color: "#607D8B", authors: ["THUDM"] },
  { color: "#3F51B5", authors: ["rednote-hilab"] },
  { color: "#DC143C", authors: ["deepseek-ai"] },
  { color: "#FF4500", authors: ["Qwen", "wan-ai","wan-ai"]},
  { color: "#009688", authors: ["XiaomiMiMo"] },
  { color: "#8BC34A", authors: ["IndexTeam"] },
];
  
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];
  const calendarKey = provider.authors[0]; // Use the same key as calendar generation
  const totalCount = calendarData[calendarKey]?.reduce((sum, day) => sum + day.count, 0) || 0;
  
  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">
      {/* 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">
        Chinese AI Labs 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 in the last year.
        </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>
      ) : (
        <>
          <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}
                />
              ))}
            </div>
          </div>

          {/* Rest of the providers */}
          {sortedProviders.length > 3 && (
              <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}
                  />
                ))}
              </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;