import React, { useState, useEffect } from "react";
import { generateCalendarData } from "../utils/calendar";
import {
  OpenSourceHeatmapProps,
  ProviderInfo,
  ModelData,
} from "../types/heatmap";
import Heatmap from "../components/Heatmap";

const PROVIDERS_MAP: Record<string, ProviderInfo> = {
  "Mistral AI": { color: "#ff7000", authors: ["mistralai"] },
  Meta: { color: "#1877F2", authors: ["facebook", "meta-llama"] },
  OpenAI: { color: "#10A37F", authors: ["openai"] },
  Anthropic: { color: "#cc785c", authors: ["Anthropic"] },
  Google: { color: "#DB4437", authors: ["google"] },
  "Allen Institute for AI": { color: "#5E35B1", authors: ["allenai"] },
  Apple: { color: "#0088cc", authors: ["apple"] },
  Microsoft: { color: "#FEB800", authors: ["microsoft"] },
  NVIDIA: { color: "#76B900", authors: ["nvidia"] },
  DeepSeek: { color: "#0088cc", authors: ["deepseek-ai"] },
  Qwen: { color: "#0088cc", authors: ["Qwen"] },
  "Cohere For AI": { color: "#4C6EE6", authors: ["CohereForAI"] },
  IBM: { color: "#4C6EE6", authors: ["ibm-granite"] },
  "Stability AI": { color: "#A020F0", authors: ["stabilityai"] },
};

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

    const entityTypes = ["models", "datasets", "spaces"];
    const allData = await Promise.all(
      uniqueAuthors.flatMap((author) =>
        entityTypes.map(async (type) => {
          const response = await fetch(
            `https://huggingface.co/api/${type}?author=${author}&sort=createdAt&direction=-1`,
          );
          const data = await response.json();
          return data.map((item: any) => ({
            createdAt: item.createdAt,
            id: item.id,
            type: type,
          }));
        }),
      ),
    );

    const flatData: ModelData[] = allData.flat();
    const calendarData = generateCalendarData(flatData, PROVIDERS_MAP);

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

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

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

  return (
    <div className="w-full max-w-screen-lg mx-auto p-4">
      <h1 className="text-3xl lg:text-5xl mt-16 font-bold text-center mb-2">
        Hugging Face Heatmap 🤗
      </h1>
      <p className="text-center text-sm my-8">
        Models, Datasets, and Spaces from the top AI labs. <br />
        Request more heatmaps by{" "}
        <a
          href="https://huggingface.co/spaces/cfahlgren1/model-release-heatmap/discussions/new"
          target="_blank"
          rel="noopener noreferrer"
          className="text-blue-500 hover:underline"
        >
          opening a discussion
        </a>
        .
      </p>
      {isLoading ? (
        <p className="text-center">Loading...</p>
      ) : (
        <div className="space-y-16">
          {Object.entries(providers)
            .sort(
              ([keyA], [keyB]) =>
                calendarData[keyB].reduce((sum, day) => sum + day.count, 0) -
                calendarData[keyA].reduce((sum, day) => sum + day.count, 0),
            )
            .map(([providerName, { color }]) => (
              <div key={providerName} className="flex flex-col items-center">
                <Heatmap
                  data={calendarData[providerName]}
                  color={color}
                  providerName={providerName}
                />
              </div>
            ))}
        </div>
      )}
    </div>
  );
};

export default OpenSourceHeatmap;