src/App.tsx

import React, { useState, useEffect, useMemo } from "react";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Input } from "@/components/ui/input";
import { mockData } from "@/lib/data";

import { ComparisonSelector } from "@/components/ComparisonSelector";
import { PricingTable } from "@/components/PricingTable";

import { BenchmarkTable } from "./components/BenchmarkTable";
import { benchmarkData } from "./lib/benchmarks/ index";
import { BenchmarkComparisonSelector } from "./components/BenchmarkComparisonSelector";
import { benchmarkMetricOrder } from "./lib/benchmarks/types";

export interface FlattenedModel extends Model {
  provider: string;
  uri: string;
  benchmark?: {
    [key: string]: number;
  };
}

export interface Model {
  name: string;
  inputPrice: number;
  outputPrice: number;
}

export interface Provider {
  provider: string;
  uri: string;
  models: Model[];
}

const App: React.FC = () => {
  const [data, setData] = useState<Provider[]>([]);
  const [comparisonModels, setComparisonModels] = useState<string[]>([]);
  const [inputTokens, setInputTokens] = useState<number>(1);
  const [outputTokens, setOutputTokens] = useState<number>(1);
  const [selectedProviders, setSelectedProviders] = useState<string[]>([]);
  const [selectedModels, setSelectedModels] = useState<string[]>([]);
  const [expandedProviders, setExpandedProviders] = useState<string[]>([]);
  const [tokenCalculation, setTokenCalculation] = useState<string>("million");
  const [benchmarkComparisonMetrics, setBenchmarkComparisonMetrics] = useState<string[]>([]);
  const [selectedBenchmarkProviders, setSelectedBenchmarkProviders] = useState<string[]>([]);
  const [selectedBenchmarkModels, setSelectedBenchmarkModels] = useState<string[]>([]);
  const [viewMode, setViewMode] = useState<"pricing" | "benchmark">("pricing");

  const [sortConfig, setSortConfig] = useState<{
    key: keyof FlattenedModel;
    direction: string;
  } | null>(null);

  const [benchmarkSortConfig, setBenchmarkSortConfig] = useState<{
    key: string;
    direction: "ascending" | "descending";
  } | null>(null);

  useEffect(() => {
    setData(mockData);
  }, []);

  const flattenDataFromPricing = (data: Provider[]): FlattenedModel[] =>
    data.flatMap((provider) =>
      provider.models.map((model) => ({
        provider: provider.provider,
        uri: provider.uri,
        ...model,
        benchmark: {},
      }))
    );

  const flattenDataFromBenchmarks = (): FlattenedModel[] =>
    benchmarkData.map((b) => ({
      provider: b.provider ?? "Unknown",
      uri: b.source,
      name: b.model,
      inputPrice: b.inputPrice,
      outputPrice: b.outputPrice,
      benchmark: b.benchmark ?? {},
    }));

  const filteredData = useMemo(() => {
    if (!selectedProviders.length && !selectedModels.length) return data;

    return data
      .filter((p) => !selectedProviders.length || selectedProviders.includes(p.provider))
      .map((p) => ({
        ...p,
        models: p.models.filter((m) => {
          if (!selectedModels.length) return selectedProviders.includes(p.provider);
          return selectedModels.includes(m.name);
        }),
      }))
      .filter((p) => p.models.length > 0);
  }, [data, selectedProviders, selectedModels]);

  const benchmarkedModels = useMemo(() => flattenDataFromBenchmarks(), []);

  const filteredBenchmarkedModels = useMemo(() => {
    return benchmarkedModels.filter((model) => {
      const providerMatch =
        selectedBenchmarkProviders.length === 0 || selectedBenchmarkProviders.includes(model.provider);
      const modelMatch =
        selectedBenchmarkModels.length === 0 || selectedBenchmarkModels.includes(model.name);

      return providerMatch && modelMatch;
    });
  }, [benchmarkedModels, selectedBenchmarkProviders, selectedBenchmarkModels]);

  const sortedBenchmarkedModels = useMemo(() => {
    if (!benchmarkSortConfig) return filteredBenchmarkedModels;

    return [...filteredBenchmarkedModels].sort((a, b) => {
      const key = benchmarkSortConfig.key;
      const isTopLevelKey = ["provider", "name", "inputPrice", "outputPrice"].includes(key);
      const aVal = isTopLevelKey ? (a as any)[key] : a.benchmark?.[key] ?? -Infinity;
      const bVal = isTopLevelKey ? (b as any)[key] : b.benchmark?.[key] ?? -Infinity;

      if (typeof aVal === "string" && typeof bVal === "string") {
        return benchmarkSortConfig.direction === "ascending"
          ? aVal.localeCompare(bVal)
          : bVal.localeCompare(aVal);
      }

      return benchmarkSortConfig.direction === "ascending"
        ? aVal - bVal
        : bVal - aVal;
    });
  }, [filteredBenchmarkedModels, benchmarkSortConfig]);

  const pricingProviders = useMemo(() => {
    const grouped: Record<string, FlattenedModel[]> = {};

    flattenDataFromPricing(data).forEach((model) => {
      const key = model.provider;
      if (!grouped[key]) grouped[key] = [];
      grouped[key].push(model);
    });

    return Object.entries(grouped).map(([provider, models]) => ({
      provider,
      uri: models[0]?.uri ?? "#",
      models: models.map(({ name, inputPrice, outputPrice }) => ({
        name,
        inputPrice,
        outputPrice,
      })),
    }));
  }, [data]);

  const benchmarkProviders = useMemo(() => {
    const grouped: Record<string, FlattenedModel[]> = {};

    benchmarkedModels.forEach((model) => {
      const key = model.provider;
      if (!grouped[key]) grouped[key] = [];
      grouped[key].push(model);
    });

    return Object.entries(grouped).map(([provider, models]) => ({
      provider,
      uri: models[0]?.uri ?? "#",
      models: models.map(({ name, inputPrice, outputPrice }) => ({
        name,
        inputPrice,
        outputPrice,
      })),
    }));
  }, [benchmarkedModels]);

  const sortedFlattenedData = useMemo(() => {
    const flattened = flattenDataFromPricing(filteredData);
    if (!sortConfig) return flattened;

    return [...flattened].sort((a, b) => {
      const aVal = a[sortConfig.key];
      const bVal = b[sortConfig.key];

      if (typeof aVal === "string" && typeof bVal === "string") {
        return sortConfig.direction === "ascending" ? aVal.localeCompare(bVal) : bVal.localeCompare(aVal);
      } else if (typeof aVal === "number" && typeof bVal === "number") {
        return sortConfig.direction === "ascending" ? aVal - bVal : bVal - aVal;
      }
      return 0;
    });
  }, [filteredData, sortConfig]);

  const requestSort = (key: keyof FlattenedModel) => {
    const direction = sortConfig?.key === key && sortConfig.direction === "ascending" ? "descending" : "ascending";
    setSortConfig({ key, direction });
  };

  const toggleProviderExpansion = (provider: string) => {
    setExpandedProviders((prev) =>
      prev.includes(provider) ? prev.filter((p) => p !== provider) : [...prev, provider]
    );
  };

  return (
    <Card className="w-full max-w-6xl mx-auto">
      <CardHeader>
        <CardTitle>LLM Pricing & Benchmark Comparison Tool</CardTitle>
      </CardHeader>
      <CardContent>
        <p className="italic text-sm text-muted-foreground mb-4">
          <a
            href="https://huggingface.co/spaces/philschmid/llm-pricing"
            className="underline"
          >
            This is a fork of philschmid tool: philschmid/llm-pricing
          </a>
        </p>

        {/* View Toggle */}
        <div className="flex justify-center gap-6 mb-6">
          <label className="flex items-center space-x-2">
            <input
              type="radio"
              name="view"
              value="pricing"
              checked={viewMode === "pricing"}
              onChange={() => setViewMode("pricing")}
            />
            <span>Pricing</span>
          </label>
          <label className="flex items-center space-x-2">
            <input
              type="radio"
              name="view"
              value="benchmark"
              checked={viewMode === "benchmark"}
              onChange={() => setViewMode("benchmark")}
            />
            <span>Benchmark</span>
          </label>
        </div>

        {viewMode === "pricing" && (
          <>
            <h3 className="text-lg font-semibold mb-2">Select Comparison Models</h3>
            <ComparisonSelector
              data={data}
              expanded={expandedProviders}
              comparisonModels={comparisonModels}
              onToggleExpand={toggleProviderExpansion}
              onChangeModel={(modelId, checked) =>
                setComparisonModels((prev) =>
                  checked ? [...prev, modelId] : prev.filter((m) => m !== modelId)
                )
              }
            />

            <div className="flex gap-4 mt-6 mb-4">
              <div className="flex-1">
                <label className="block text-sm font-medium">Input Tokens ({tokenCalculation})</label>
                <Input type="number" value={inputTokens} min={1} onChange={(e) => setInputTokens(Number(e.target.value))} />
              </div>
              <div className="flex-1">
                <label className="block text-sm font-medium">Output Tokens ({tokenCalculation})</label>
                <Input type="number" value={outputTokens} min={1} onChange={(e) => setOutputTokens(Number(e.target.value))} />
              </div>
              <div className="flex-1">
                <label className="block text-sm font-medium">Token Calculation</label>
                <select
                  value={tokenCalculation}
                  onChange={(e) => setTokenCalculation(e.target.value)}
                  className="mt-1 block w-full pl-3 pr-10 py-2 text-base border rounded-md"
                >
                  <option value="billion">Billion Tokens</option>
                  <option value="million">Million Tokens</option>
                  <option value="thousand">Thousand Tokens</option>
                  <option value="unit">Unit Tokens</option>
                </select>
              </div>
            </div>

            <h2 className="text-lg font-semibold mb-2">Pricing Table</h2>
            <PricingTable
              data={sortedFlattenedData}
              providers={pricingProviders}
              selectedProviders={selectedProviders}
              selectedModels={selectedModels}
              onProviderChange={setSelectedProviders}
              onModelChange={setSelectedModels}
              comparisonModels={comparisonModels}
              inputTokens={inputTokens}
              outputTokens={outputTokens}
              tokenCalculation={tokenCalculation}
              requestSort={requestSort}
              sortConfig={sortConfig}
            />
          </>
        )}

        {viewMode === "benchmark" && (
          <>
            <h3 className="text-lg font-semibold mb-2">Select Benchmark Metrics to Compare</h3>
            <BenchmarkComparisonSelector
              allMetrics={benchmarkMetricOrder.filter(
                (metric) => benchmarkedModels.some((m) => m.benchmark?.[metric] !== undefined)
              )}
              selected={benchmarkComparisonMetrics}
              onChange={(metric, checked) =>
                setBenchmarkComparisonMetrics((prev) =>
                  checked ? [...prev, metric] : prev.filter((m) => m !== metric)
                )
              }
            />

            <h2 className="text-lg font-semibold mb-2">Benchmark Table</h2>
            <BenchmarkTable
              data={sortedBenchmarkedModels}
              comparisonMetrics={benchmarkComparisonMetrics}
              requestSort={(key) => {
                setBenchmarkSortConfig((prev) =>
                  prev?.key === key
                    ? { key, direction: prev.direction === "ascending" ? "descending" : "ascending" }
                    : { key, direction: "descending" }
                );
              }}
              sortConfig={benchmarkSortConfig}
              allProviders={benchmarkProviders}
              selectedProviders={selectedBenchmarkProviders}
              selectedModels={selectedBenchmarkModels}
              onProviderChange={setSelectedBenchmarkProviders}
              onModelChange={setSelectedBenchmarkModels}
            />
          </>
        )}
      </CardContent>
    </Card>
  );
};

export default App;