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: "#0066CC", authors: ["apple"] },
  { color: "#FEB800", authors: ["microsoft"] },
  { color: "#76B900", authors: ["nvidia"] },
  { color: "#00A8E0", authors: ["deepseek-ai"] },
  { color: "#6366F1", authors: ["Qwen"] },
  { color: "#FF6B6B", authors: ["CohereLabs"] },
  { color: "#4ECDC4", authors: ["ibm-granite"] },
  { color: "#A855F7", 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,
    };
  } catch (error) {
    console.error("Error fetching data:", error);
    return {
      props: {
        calendarData: {},
        providers: PROVIDERS,
      },
      revalidate: 60,
    };
  }
}

const ProviderCard = React.memo(({ 
  provider, 
  calendarData, 
  isExpanded, 
  onToggle 
}: { 
  provider: ProviderInfo, 
  calendarData: CalendarData,
  isExpanded: boolean,
  onToggle: () => void
}) => {
  const providerName = provider.fullName || provider.authors[0];
  const totalActivity = calendarData[providerName]?.reduce((sum, day) => sum + day.count, 0) || 0;
  const recentActivity = calendarData[providerName]?.slice(-30).reduce((sum, day) => sum + day.count, 0) || 0;

  return (
    <div className="bg-white rounded-xl shadow-sm border border-gray-200 overflow-hidden transition-all duration-200 hover:shadow-md">
      <button
        onClick={onToggle}
        className="w-full px-6 py-4 flex items-center justify-between hover:bg-gray-50 transition-colors"
      >
        <div className="flex items-center gap-4">
          {provider.avatarUrl && (
            <img 
              src={provider.avatarUrl} 
              alt={providerName}
              className="w-12 h-12 rounded-lg object-cover"
            />
          )}
          <div className="text-left">
            <h3 className="font-semibold text-lg text-gray-900">{providerName}</h3>
            <div className="flex gap-4 text-sm text-gray-600">
              <span>Total: <span className="font-medium text-gray-900">{totalActivity.toLocaleString()}</span></span>
              <span>Last 30 days: <span className="font-medium text-gray-900">{recentActivity.toLocaleString()}</span></span>
            </div>
          </div>
        </div>
        <div className="flex items-center gap-3">
          <div className="w-3 h-3 rounded-full" style={{ backgroundColor: provider.color }}></div>
          <svg 
            className={`w-5 h-5 text-gray-400 transition-transform ${isExpanded ? 'rotate-180' : ''}`}
            fill="none" 
            stroke="currentColor" 
            viewBox="0 0 24 24"
          >
            <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M19 9l-7 7-7-7" />
          </svg>
        </div>
      </button>
      
      {isExpanded && (
        <div className="px-6 pb-6 border-t border-gray-100">
          <div className="pt-4">
            <Heatmap
              data={calendarData[providerName]}
              color={provider.color}
              providerName={providerName}
              fullName={provider.fullName ?? providerName}
              avatarUrl={provider.avatarUrl ?? ''}
            />
          </div>
        </div>
      )}
    </div>
  );
});

const OpenSourceHeatmap: React.FC<OpenSourceHeatmapProps> = ({
  calendarData,
  providers,
}) => {
  const [isLoading, setIsLoading] = useState(true);
  const [expandedProviders, setExpandedProviders] = useState<Set<string>>(new Set());
  const [filterQuery, setFilterQuery] = useState("");
  const [sortBy, setSortBy] = useState<"total" | "recent">("total");

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

  const sortedAndFilteredProviders = useMemo(() => {
    let filtered = providers;
    
    if (filterQuery) {
      filtered = providers.filter(p => 
        (p.fullName || p.authors[0]).toLowerCase().includes(filterQuery.toLowerCase())
      );
    }

    return filtered.sort((a, b) => {
      const aName = a.fullName || a.authors[0];
      const bName = b.fullName || b.authors[0];
      
      if (sortBy === "total") {
        return calendarData[bName]?.reduce((sum, day) => sum + day.count, 0) || 0 -
               calendarData[aName]?.reduce((sum, day) => sum + day.count, 0) || 0;
      } else {
        return calendarData[bName]?.slice(-30).reduce((sum, day) => sum + day.count, 0) || 0 -
               calendarData[aName]?.slice(-30).reduce((sum, day) => sum + day.count, 0) || 0;
      }
    });
  }, [providers, calendarData, filterQuery, sortBy]);

  const toggleProvider = (providerName: string) => {
    setExpandedProviders(prev => {
      const newSet = new Set(prev);
      if (newSet.has(providerName)) {
        newSet.delete(providerName);
      } else {
        newSet.add(providerName);
      }
      return newSet;
    });
  };

  const toggleAll = () => {
    if (expandedProviders.size === sortedAndFilteredProviders.length) {
      setExpandedProviders(new Set());
    } else {
      setExpandedProviders(new Set(sortedAndFilteredProviders.map(p => p.fullName || p.authors[0])));
    }
  };

  return (
    <div className="min-h-screen bg-gray-50">
      <div className="w-full max-w-7xl mx-auto px-4 py-12">
        {/* Header */}
        <div className="text-center mb-12">
          <h1 className="text-5xl font-bold text-gray-900 mb-4">
            AI Labs Activity Dashboard
          </h1>
          <p className="text-lg text-gray-600 mb-8">
            Track models, datasets, and spaces from leading AI organizations on Hugging Face
          </p>
          
          {/* Controls */}
          <div className="flex flex-col sm:flex-row gap-4 justify-center items-center max-w-2xl mx-auto">
            <div className="relative flex-1 w-full">
              <input
                type="text"
                placeholder="Search organizations..."
                value={filterQuery}
                onChange={(e) => setFilterQuery(e.target.value)}
                className="w-full px-4 py-2 pl-10 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
              />
              <svg className="absolute left-3 top-2.5 w-5 h-5 text-gray-400" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z" />
              </svg>
            </div>
            
            <div className="flex gap-2">
              <select
                value={sortBy}
                onChange={(e) => setSortBy(e.target.value as "total" | "recent")}
                className="px-4 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent"
              >
                <option value="total">Sort by Total</option>
                <option value="recent">Sort by Recent</option>
              </select>
              
              <button
                onClick={toggleAll}
                className="px-4 py-2 bg-gray-100 hover:bg-gray-200 rounded-lg transition-colors"
              >
                {expandedProviders.size === sortedAndFilteredProviders.length ? 'Collapse All' : 'Expand All'}
              </button>
              
              <UserSearchDialog />
            </div>
          </div>
        </div>

        {/* Stats Summary */}
        {!isLoading && (
          <div className="grid grid-cols-1 sm:grid-cols-3 gap-4 mb-8">
            <div className="bg-white rounded-lg p-6 text-center shadow-sm border border-gray-200">
              <div className="text-3xl font-bold text-gray-900">
                {providers.length}
              </div>
              <div className="text-sm text-gray-600">Organizations</div>
            </div>
            <div className="bg-white rounded-lg p-6 text-center shadow-sm border border-gray-200">
              <div className="text-3xl font-bold text-gray-900">
                {Object.values(calendarData).reduce((total, days) => 
                  total + days.reduce((sum, day) => sum + day.count, 0), 0
                ).toLocaleString()}
              </div>
              <div className="text-sm text-gray-600">Total Activities</div>
            </div>
            <div className="bg-white rounded-lg p-6 text-center shadow-sm border border-gray-200">
              <div className="text-3xl font-bold text-gray-900">
                {Object.values(calendarData).reduce((total, days) => 
                  total + days.slice(-30).reduce((sum, day) => sum + day.count, 0), 0
                ).toLocaleString()}
              </div>
              <div className="text-sm text-gray-600">Last 30 Days</div>
            </div>
          </div>
        )}

        {/* Provider List */}
        {isLoading ? (
          <div className="flex justify-center items-center h-64">
            <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-gray-900"></div>
          </div>
        ) : (
          <div className="space-y-4">
            {sortedAndFilteredProviders.length === 0 ? (
              <div className="text-center py-12 text-gray-500">
                No organizations found matching "{filterQuery}"
              </div>
            ) : (
              sortedAndFilteredProviders.map((provider) => {
                const providerName = provider.fullName || provider.authors[0];
                return (
                  <ProviderCard
                    key={providerName}
                    provider={provider}
                    calendarData={calendarData}
                    isExpanded={expandedProviders.has(providerName)}
                    onToggle={() => toggleProvider(providerName)}
                  />
                );
              })
            )}
          </div>
        )}
      </div>
    </div>
  );
};

export default React.memo(OpenSourceHeatmap);