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 [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={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 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;