Dashboard/components/dashboard/charts/model-distribution-chart.tsx

"use client";

import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import {
  BarChart,
  Bar,
  XAxis,
  YAxis,
  CartesianGrid,
  Tooltip,
  ResponsiveContainer,
} from "recharts";

interface ModelDistributionChartProps {
  title: string;
  subtitle?: string;
  data: Array<{
    name: string;
    value: number;
    color?: string;
    models?: Array<{
      name: string;
      value: number;
    }>;
  }>;
  showLegend?: boolean;
  totalTokens?: number;
}

interface ModelData {
  name: string;
  value: number;
}

interface StackedDataEntry {
  provider: string;
  total: number;
  models: ModelData[];
  baseColor: string;
  modelColors: string[];
  [key: string]: string | number | ModelData[] | string[];
}

interface ModelInfo {
  color: string;
  name: string;
  provider: string;
}

// Couleurs par fournisseur (couleurs de base)
const providerColors: { [key: string]: string } = {
  Anthropic: "#7C3AED", // Violet vif
  OpenAI: "#059669", // Vert turquoise vif
  "Mistral AI": "#D97706", // Orange vif
  Meta: "#DB2777", // Rose/Magenta vif
  Google: "#2563EB", // Bleu vif
  Cohere: "#0891B2", // Cyan vif
};

// Fonction pour générer des variations de couleur pour les modèles d'un même provider
const generateModelColors = (baseColor: string, modelCount: number) => {
  const colors = [];
  for (let i = 0; i < modelCount; i++) {
    // Créer des variations en ajustant la luminosité
    const opacity = 1 - i * 0.15; // De 1.0 à 0.4 environ
    colors.push(
      `${baseColor}${Math.round(opacity * 255)
        .toString(16)
        .padStart(2, "0")}`
    );
  }
  return colors;
};

// Fonction pour regrouper les modèles par fournisseur et préparer les données pour le graphique empilé
const prepareStackedData = (
  modelData: Array<{ name: string; value: number }>
) => {
  const providerMap: {
    [key: string]: {
      value: number;
      models: Array<{ name: string; value: number }>;
    };
  } = {};

  modelData.forEach((model) => {
    let provider = "";

    // Déterminer le fournisseur basé sur le nom du modèle
    if (
      model.name.toLowerCase().includes("claude") ||
      model.name.toLowerCase().includes("anthropic")
    ) {
      provider = "Anthropic";
    } else if (
      model.name.toLowerCase().includes("gpt") ||
      model.name.toLowerCase().includes("openai")
    ) {
      provider = "OpenAI";
    } else if (model.name.toLowerCase().includes("mistral")) {
      provider = "Mistral AI";
    } else if (
      model.name.toLowerCase().includes("llama") ||
      model.name.toLowerCase().includes("meta")
    ) {
      provider = "Meta";
    } else if (
      model.name.toLowerCase().includes("palm") ||
      model.name.toLowerCase().includes("gemini") ||
      model.name.toLowerCase().includes("google")
    ) {
      provider = "Google";
    } else if (model.name.toLowerCase().includes("cohere")) {
      provider = "Cohere";
    } else {
      provider = "Autres";
    }

    if (!providerMap[provider]) {
      providerMap[provider] = { value: 0, models: [] };
    }

    providerMap[provider].value += model.value;
    providerMap[provider].models.push(model);
  });

  // Créer les données pour le graphique empilé
  const stackedData: StackedDataEntry[] = Object.entries(providerMap).map(
    ([providerName, data]) => {
      const baseColor = providerColors[providerName] || "#6B7280";
      const modelColors = generateModelColors(baseColor, data.models.length);

      // Créer un objet avec le provider comme clé et chaque modèle comme propriété
      const stackedEntry: StackedDataEntry = {
        provider: providerName,
        total: data.value,
        models: data.models,
        baseColor,
        modelColors,
      };

      // Ajouter chaque modèle comme propriété séparée pour le stacking
      data.models.forEach((model, index) => {
        const modelKey = `${providerName}_${model.name}`;
        stackedEntry[modelKey] = model.value;
        stackedEntry[`${modelKey}_color`] = modelColors[index];
        stackedEntry[`${modelKey}_name`] = model.name;
      });

      return stackedEntry;
    }
  );

  // Créer la liste de tous les modèles uniques pour les barres
  const allModelKeys: string[] = [];
  const modelInfo: { [key: string]: ModelInfo } = {};

  stackedData.forEach((entry) => {
    entry.models.forEach((model, index) => {
      const modelKey = `${entry.provider}_${model.name}`;
      allModelKeys.push(modelKey);
      modelInfo[modelKey] = {
        color: entry.modelColors[index],
        name: model.name,
        provider: entry.provider,
      };
    });
  });

  return { stackedData, allModelKeys, modelInfo };
};

interface CustomStackedTooltipProps {
  active?: boolean;
  payload?: Array<{
    payload: StackedDataEntry;
  }>;
  label?: string;
}

const CustomStackedTooltip = ({
  active,
  payload,
  label,
}: CustomStackedTooltipProps) => {
  if (active && payload && payload.length) {
    const providerData = payload[0].payload;
    const totalTokens = providerData.total;

    return (
      <div className="bg-white p-3 border rounded-lg shadow-lg">
        <p className="font-semibold">{label}</p>
        <p className="text-blue-600 font-medium">
          Total: {totalTokens.toLocaleString()} tokens
        </p>
        {providerData.models && (
          <div className="mt-2 space-y-1">
            <p className="text-sm font-medium text-gray-600">Modèles:</p>
            {providerData.models.map((model, index) => (
              <div key={index} className="flex items-center gap-2">
                <div
                  className="w-3 h-3 rounded-full"
                  style={{ backgroundColor: providerData.modelColors[index] }}
                />
                <p className="text-xs text-gray-700">
                  {model.name}: {model.value.toLocaleString()} tokens
                </p>
              </div>
            ))}
          </div>
        )}
      </div>
    );
  }
  return null;
};

export function ModelDistributionChart({
  title,
  subtitle,
  data,
  totalTokens,
}: ModelDistributionChartProps) {
  // Si les données sont déjà groupées par fournisseur, les utiliser directement
  // Sinon, les regrouper automatiquement
  const modelData = data[0]?.models
    ? data.flatMap((d) => d.models || [])
    : data;
  const { stackedData, allModelKeys, modelInfo } =
    prepareStackedData(modelData);

  return (
    <Card>
      <CardHeader className="pb-2">
        <CardTitle className="text-sm font-medium text-muted-foreground">
          {title}
        </CardTitle>
        {subtitle && (
          <p className="text-xs text-muted-foreground mt-1">{subtitle}</p>
        )}
      </CardHeader>
      <CardContent className="pt-0">
        <ResponsiveContainer width="100%" height={300}>
          <BarChart
            data={stackedData}
            margin={{ top: 10, right: 10, left: 10, bottom: 60 }}
          >
            <CartesianGrid strokeDasharray="3 3" className="stroke-muted/20" />
            <XAxis
              dataKey="provider"
              axisLine={false}
              tickLine={false}
              className="text-xs fill-muted-foreground"
              angle={-45}
              textAnchor="end"
              height={80}
              interval={0}
            />
            <YAxis
              axisLine={false}
              tickLine={false}
              className="text-xs fill-muted-foreground"
              tickFormatter={(value) => {
                if (value >= 1000000) return `${(value / 1000000).toFixed(1)}M`;
                if (value >= 1000) return `${(value / 1000).toFixed(1)}K`;
                return value.toString();
              }}
            />
            <Tooltip content={<CustomStackedTooltip />} />

            {/* Créer une barre empilée pour chaque modèle */}
            {allModelKeys.map((modelKey) => (
              <Bar
                key={modelKey}
                dataKey={modelKey}
                stackId="models"
                fill={modelInfo[modelKey].color}
                radius={0}
              />
            ))}
          </BarChart>
        </ResponsiveContainer>

        {/* Légende des providers avec leurs couleurs de base */}
        <div className="mt-4 grid grid-cols-2 gap-3">
          {stackedData.map((item, index) => (
            <div
              key={index}
              className="p-3 rounded-lg border border-muted/20 bg-muted/5 hover:bg-muted/10 transition-colors"
              style={{
                borderLeftColor: item.baseColor,
                borderLeftWidth: "3px",
                borderLeftStyle: "solid",
              }}
            >
              <div className="flex items-center gap-2 mb-1">
                <div
                  className="w-2 h-2 rounded-full"
                  style={{ backgroundColor: item.baseColor }}
                ></div>
                <h3
                  className="text-sm font-medium"
                  style={{ color: item.baseColor }}
                >
                  {item.provider}
                </h3>
              </div>
              <p className="text-lg font-semibold text-foreground">
                {item.total.toLocaleString()}
              </p>
              <p className="text-xs text-muted-foreground">
                {item.models.length} modèle{item.models.length > 1 ? "s" : ""}
              </p>
            </div>
          ))}
        </div>

        {/* Total général */}
        {totalTokens && (
          <div className="mt-4 pt-3 border-t border-muted/20 text-center">
            <p className="text-sm text-muted-foreground">
              Total général:{" "}
              <span className="font-semibold text-foreground">
                {totalTokens.toLocaleString()}
              </span>{" "}
              tokens
            </p>
          </div>
        )}

        {/* Légende détaillée des modèles avec leurs couleurs spécifiques */}
        <div className="mt-4 pt-3 border-t border-muted/20">
          <h4 className="text-sm font-medium text-muted-foreground mb-3 text-center">
            Modèles par provider
          </h4>
          <div className="space-y-3">
            {stackedData.map((provider) => (
              <div key={provider.provider} className="space-y-1">
                <h5
                  className="text-xs font-medium"
                  style={{ color: provider.baseColor }}
                >
                  {provider.provider}
                </h5>
                <div className="flex flex-wrap gap-x-4 gap-y-1 ml-2">
                  {provider.models.map((model, index) => (
                    <div key={index} className="flex items-center gap-1">
                      <div
                        className="w-2 h-2 rounded-full flex-shrink-0"
                        style={{ backgroundColor: provider.modelColors[index] }}
                      ></div>
                      <span className="text-xs text-muted-foreground">
                        {model.name} ({model.value.toLocaleString()})
                      </span>
                    </div>
                  ))}
                </div>
              </div>
            ))}
          </div>
        </div>
      </CardContent>
    </Card>
  );
}