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";


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 [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, // placeholder or use an optional `b.link` if available
      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);
          if (!selectedModels.length) return !selectedProviders.length || selectedProviders.includes(p.provider);
          return selectedModels.includes(m.name);
        }),
      }))
      .filter((p) => p.models.length > 0);
  }, [data, selectedProviders, selectedModels]);

  const benchmarkedModels = useMemo(() => {
    return 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);

      // When not linking, allow filtering by either
      if (selectedBenchmarkProviders.length > 0 && selectedBenchmarkModels.length > 0) {
        return providerMatch || modelMatch;
      }

      return providerMatch && modelMatch; // this handles the case where one or both are empty
    });
  }, [
    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 Calculator</CardTitle>

      </CardHeader>

      <CardContent>

        {/* Source Link */}

        <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>



        {/* Comparison Model Selector */}

        <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)

            )

          }

        />



        {/* Token Inputs */}

        <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>



        {/* Pricing Table */}

        <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}

        />





        {/* Benchmark Table */}

        <h3 className="text-lg font-semibold mt-12 mb-2">Select Benchmark Metrics to Compare</h3>

        <BenchmarkComparisonSelector

          allMetrics={Array.from(new Set(benchmarkedModels.flatMap((m) => Object.keys(m.benchmark ?? {})))).sort()}

          selected={benchmarkComparisonMetrics}

          onChange={(metric, checked) =>

            setBenchmarkComparisonMetrics((prev) =>

              checked ? [...prev, metric] : prev.filter((m) => m !== metric)

            )

          }

        />



        <h2 className="text-lg font-semibold mt-12 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;