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

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 [linkProviderModel, setLinkProviderModel] = useState<boolean>(false);
  const [benchmarkComparisonMetrics, setBenchmarkComparisonMetrics] = useState<string[]>([]);
  const [selectedBenchmarkProviders, setSelectedBenchmarkProviders] = useState<string[]>([]);
  const [selectedBenchmarkModels, setSelectedBenchmarkModels] = useState<string[]>([]);
  const [linkBenchmarkProviderModel, setLinkBenchmarkProviderModel] = useState<boolean>(false);


  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 && !linkProviderModel) return data;

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

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

    if (linkBenchmarkProviderModel) {
      return providerMatch && modelMatch;
    }

    // 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,
  linkBenchmarkProviderModel,
]);


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

    return [...benchmarkedModels].sort((a, b) => {
      const key = benchmarkSortConfig.key;
      const aVal = key === "provider" || key === "name"
        ? (a as any)[key]?.toLowerCase?.() ?? ""
        : a.benchmark?.[key] ?? -Infinity;
      const bVal = key === "provider" || key === "name"
        ? (b as any)[key]?.toLowerCase?.() ?? ""
        : 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 benchmarkProviders = useMemo(() => {
    const modelsInBenchmark = new Set(
      benchmarkedModels.map((m) => `${m.provider}:${m.name}`)
    );

    return data
      .map((provider) => ({
        ...provider,
        models: provider.models.filter((model) =>
          modelsInBenchmark.has(`${provider.provider}:${model.name}`)
        ),
      }))
      .filter((p) => p.models.length > 0);
  }, [data, 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>

        {/* Provider–Model Toggle */}
        <div className="flex items-center space-x-2 mb-6">
          <Switch id="linkProviderModel" checked={linkProviderModel} onCheckedChange={setLinkProviderModel} />
          <label htmlFor="linkProviderModel" className="text-sm">Link Provider and Model</label>
        </div>

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


        {/* 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}
          linkProviderModel={linkBenchmarkProviderModel}
          onLinkToggle={setLinkBenchmarkProviderModel}
        />


      </CardContent>
    </Card>
  );
};

export default App;