coai/app/src/components/admin/assemblies/ModelUsageChart.tsx

import { useMemo } from "react";
import { useTranslation } from "react-i18next";
import { Loader2 } from "lucide-react";
import Tips from "@/components/Tips.tsx";
import { sum } from "@/utils/base.ts";
import { DonutChart, Legend } from "@tremor/react";
import { getReadableNumber } from "@/utils/processor.ts";
import { getModelColor } from "@/admin/colors.ts";

type ModelChartProps = {
  labels: string[];
  datasets: {
    model: string;
    data: number[];
  }[];
};

type DataUsage = {
  name: string;
  value: number;
};

function ModelUsageChart({ labels, datasets }: ModelChartProps) {
  const { t } = useTranslation();

  const usage = useMemo((): Record<string, number> => {
    const usage: Record<string, number> = {};
    datasets.forEach((dataset) => {
      usage[dataset.model] = sum(dataset.data);
    });
    return usage;
  }, [datasets]);

  const data = useMemo((): DataUsage[] => {
    const models: string[] = Object.keys(usage);
    const data: number[] = models.map((model) => usage[model]);

    return models.map(
      (model, i): DataUsage => ({ name: model, value: data[i] }),
    );
  }, [usage]);

  const sorted = useMemo(() => {
    return data.sort((a, b) => b.value - a.value);
  }, [data]);

  const categories = useMemo(() => {
    return sorted.map(
      (item) => `${item.name} (${getReadableNumber(item.value, 1)})`,
    );
  }, [sorted]);

  type CustomTooltipTypeDonut = {
    payload: any;
    active: boolean | undefined;
    label: any;
  };

  const customTooltip = (props: CustomTooltipTypeDonut) => {
    const { payload, active } = props;
    if (!active || !payload) return null;
    const categoryPayload = payload?.[0];
    if (!categoryPayload) return null;
    return (
      <div className="chart-tooltip min-w-56 w-max z-10 rounded-tremor-default border border-tremor-border bg-tremor-background p-2 text-tremor-default shadow-tremor-dropdown">
        <div className="flex flex-1 space-x-2.5">
          <div
            className={`flex w-1.5 flex-col bg-${categoryPayload?.color} rounded`}
          />
          <div className="w-full">
            <div className="flex items-center justify-between space-x-8">
              <p className="whitespace-nowrap text-right text-tremor-content">
                {categoryPayload.name}
              </p>
              <p className="whitespace-nowrap text-right font-medium text-tremor-content-emphasis">
                {getReadableNumber(categoryPayload.value, 1)} tokens
              </p>
            </div>
          </div>
        </div>
      </div>
    );
  };

  return (
    <div className={`chart`}>
      <div className={`chart-title mb-2`}>
        <div className={`flex flex-row items-center`}>
          {t("admin.model-usage-chart")}
          <Tips content={t("admin.model-chart-tip")} />
        </div>
        {labels.length === 0 && (
          <Loader2 className={`h-4 w-4 inline-block animate-spin`} />
        )}
      </div>
      <div className={`flex flex-row`}>
        <DonutChart
          className={`common-chart p-4 w-[50%]`}
          variant={`donut`}
          data={data}
          showAnimation={true}
          valueFormatter={(value) => getReadableNumber(value, 1)}
          customTooltip={customTooltip}
          colors={data.map((item) => getModelColor(item.name))}
        />
        <Legend
          className={`common-chart p-2 w-[50%] z-0`}
          // keep 6 items max
          categories={categories.slice(0, 6)}
          colors={sorted.slice(0, 6).map((item) => getModelColor(item.name))}
        />
      </div>
    </div>
  );
}

export default ModelUsageChart;