import React, { useState, useEffect } from "react";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Checkbox } from "@/components/ui/checkbox";
import { Input } from "@/components/ui/input";
import {
  Table,
  TableBody,
  TableCell,
  TableHead,
  TableHeader,
  TableRow,
} from "@/components/ui/table";
import { MultiSelect } from "@/components/ui/multi-select";
import {
  Collapsible,
  CollapsibleContent,
  CollapsibleTrigger,
} from "@/components/ui/collapsible";
import { Button } from "@/components/ui/button";
import { ChevronDown, ChevronRight } from "lucide-react";
import { mockData } from "./lib/data"; // Assuming you have this file for mock data
import { Switch } from "@/components/ui/switch";

interface FlattenedModel extends Model {
  provider: string;
  uri: string;
}

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 [sortConfig, setSortConfig] = useState<{
    key: keyof FlattenedModel;
    direction: string;
  } | null>(null);

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

  const calculatePrice = (price: number, tokens: number): number => {
    let multiplier = 1;
    if (tokenCalculation === "thousand") {
      multiplier = 1e-3;
    } else if (tokenCalculation === "unit") {
      multiplier = 1e-6;
    } else if (tokenCalculation === "billion") {
      multiplier = 1e3;
    }
    return price * tokens * multiplier;
  };

  const calculateComparison = (
    modelPrice: number,
    comparisonPrice: number
  ): string => {
    return (((modelPrice - comparisonPrice) / comparisonPrice) * 100).toFixed(
      2
    );
  };

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

  const filteredData =
    selectedProviders.length === 0 &&
    selectedModels.length === 0 &&
    !linkProviderModel
      ? data.map((provider) => ({
          ...provider,
          models: provider.models,
        }))
      : data
          .filter(
            (provider) =>
              selectedProviders.length === 0 ||
              selectedProviders.includes(provider.provider)
          )
          .map((provider) => ({
            ...provider,
            models: provider.models.filter((model) => {
              // If linking is enabled and no models are selected, filter by provider
              if (linkProviderModel && selectedModels.length === 0)
                return selectedProviders.includes(provider.provider);

              // If no models are selected and linking is off, show all models from selected providers (or all if no providers selected)
              if (!linkProviderModel && selectedModels.length === 0)
                return (
                  selectedProviders.length === 0 ||
                  selectedProviders.includes(provider.provider)
                );

              // Otherwise, only show selected models
              return selectedModels.includes(model.name);
            }),
          }))
          .filter((provider) => provider.models.length > 0);

  const sortedFlattenedData = React.useMemo(() => {
    let sortableData: FlattenedModel[] = flattenData(filteredData);
    if (sortConfig !== null) {
      sortableData.sort((a, b) => {
        const aValue = a[sortConfig.key];
        const bValue = b[sortConfig.key];

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

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

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

  const getModelsForSelectedProviders = () => {
    if (!linkProviderModel) {
      return data
        .flatMap((provider) =>
          provider.models.map((model) => ({
            label: model.name,
            value: model.name,
            provider: provider.provider,
          }))
        )
        .reduce(
          (
            acc: { label: string; value: string; provider: string }[],
            curr: { label: string; value: string; provider: string }
          ) => {
            if (!acc.find((m) => m.value === curr.value)) {
              acc.push(curr);
            }
            return acc;
          },
          []
        );
    }

    return data
      .filter((provider) => selectedProviders.includes(provider.provider))
      .flatMap((provider) =>
        provider.models.map((model) => ({
          label: model.name,
          value: model.name,
          provider: provider.provider,
        }))
      )
      .reduce(
        (
          acc: { label: string; value: string; provider: string }[],
          curr: { label: string; value: string; provider: string }
        ) => {
          if (!acc.find((m) => m.value === curr.value)) {
            acc.push(curr);
          }
          return acc;
        },
        []
      );
  };

  return (
    <Card className="w-full max-w-6xl mx-auto">
      <CardHeader>
        <CardTitle>LLM Pricing Calculator</CardTitle>
      </CardHeader>
      <CardContent>
        <div className="mb-4">
          <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>
          <h3 className="text-lg font-semibold mb-2">
            Select Comparison Models
          </h3>
          <div className="grid grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4">
            {data.map((provider) => (
              <Collapsible
                key={provider.provider}
                open={expandedProviders.includes(provider.provider)}
                onOpenChange={() => toggleProviderExpansion(provider.provider)}
              >
                <CollapsibleTrigger asChild>
                  <Button variant="outline" className="w-full justify-between">
                    {provider.provider}
                    {expandedProviders.includes(provider.provider) ? (
                      <ChevronDown className="h-4 w-4" />
                    ) : (
                      <ChevronRight className="h-4 w-4" />
                    )}
                  </Button>
                </CollapsibleTrigger>
                <CollapsibleContent className="mt-2">
                  {provider.models.map((model) => (
                    <div
                      key={`${provider.provider}:${model.name}`}
                      className="flex items-center space-x-2 mb-1"
                    >
                      <Checkbox
                        id={`${provider.provider}:${model.name}`}
                        checked={comparisonModels.includes(
                          `${provider.provider}:${model.name}`
                        )}
                        onCheckedChange={(checked) => {
                          if (checked) {
                            setComparisonModels((prev) => [
                              ...prev,
                              `${provider.provider}:${model.name}`,
                            ]);
                          } else {
                            setComparisonModels((prev) =>
                              prev.filter(
                                (m) =>
                                  m !== `${provider.provider}:${model.name}`
                              )
                            );
                          }
                        }}
                      />
                      <label
                        htmlFor={`${provider.provider}:${model.name}`}
                        className="text-sm font-medium text-gray-700"
                      >
                        {model.name}
                      </label>
                    </div>
                  ))}
                </CollapsibleContent>
              </Collapsible>
            ))}
          </div>
        </div>

        <div className="flex gap-4 mb-4">
          <div className="flex-1">
            <label
              htmlFor="inputTokens"
              className="block text-sm font-medium text-gray-700"
            >
              Input Tokens ({tokenCalculation})
            </label>
            <Input
              id="inputTokens"
              type="number"
              value={inputTokens}
              min={1}
              onChange={(e) => setInputTokens(Number(e.target.value))}
              className="mt-1"
            />
          </div>
          <div className="flex-1">
            <label
              htmlFor="outputTokens"
              className="block text-sm font-medium text-gray-700"
            >
              Output Tokens ({tokenCalculation})
            </label>
            <Input
              id="outputTokens"
              type="number"
              value={outputTokens}
              min={1}
              onChange={(e) => setOutputTokens(Number(e.target.value))}
              className="mt-1"
            />
          </div>
          <div className="flex-1">
            <label
              htmlFor="tokenCalculation"
              className="block text-sm font-medium text-gray-700"
            >
              Token Calculation
            </label>
            <select
              id="tokenCalculation"
              value={tokenCalculation}
              onChange={(e) => setTokenCalculation(e.target.value)}
              className="mt-1 block w-full pl-3 pr-10 py-2 text-base bg-white border focus:outline-none focus:ring-indigo-500  sm:text-sm 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>

        <p className="italic text-sm text-muted-foreground mb-4">
          Note: If you use Amazon Bedrock or Azure prices for Anthropic, Cohere
          or OpenAI should be the same.
        </p>
        <div className="flex items-center space-x-2 mb-4">
          <Switch
            id="linkProviderModel"
            checked={linkProviderModel}
            onCheckedChange={setLinkProviderModel}
          />
          <label htmlFor="linkProviderModel" className="text-sm">
            Link Provider and Model
          </label>
        </div>

        <Table>
          <TableHeader>
            <TableRow>
              <TableHead>
                <button type="button" onClick={() => requestSort("provider")}>
                  Provider{" "}
                  {sortConfig?.key === "provider"
                    ? sortConfig.direction === "ascending"
                      ? "▲"
                      : "▼"
                    : null}
                </button>
              </TableHead>
              <TableHead>
                <button type="button" onClick={() => requestSort("name")}>
                  Model{" "}
                  {sortConfig?.key === "name"
                    ? sortConfig.direction === "ascending"
                      ? "▲"
                      : "▼"
                    : null}
                </button>
              </TableHead>

              <TableHead>
                <button type="button" onClick={() => requestSort("inputPrice")}>
                  Input Price (million tokens)
                  {sortConfig?.key === "inputPrice"
                    ? sortConfig.direction === "ascending"
                      ? "▲"
                      : "▼"
                    : null}
                </button>
              </TableHead>
              <TableHead>
                <button
                  type="button"
                  onClick={() => requestSort("outputPrice")}
                >
                  Output Price (million tokens)
                  {sortConfig?.key === "outputPrice"
                    ? sortConfig.direction === "ascending"
                      ? "▲"
                      : "▼"
                    : null}
                </button>
              </TableHead>

              <TableHead>
                Total Price (per {tokenCalculation} tokens){" "}
              </TableHead>
              {comparisonModels.map((model) => (
                <TableHead key={model} colSpan={2}>
                  Compared to {model}
                </TableHead>
              ))}
            </TableRow>
            <TableRow>
              <TableHead>
                <MultiSelect
                  options={
                    data.map((provider) => ({
                      label: provider.provider,
                      value: provider.provider,
                    })) || []
                  }
                  onValueChange={setSelectedProviders}
                  defaultValue={selectedProviders}
                />
              </TableHead>
              <TableHead>
                <MultiSelect
                  options={getModelsForSelectedProviders()}
                  defaultValue={selectedModels}
                  onValueChange={setSelectedModels}
                />
              </TableHead>
              <TableHead />
              <TableHead />
              <TableHead />
              {comparisonModels.flatMap((model) => [
                <TableHead key={`${model}-input`}>Input</TableHead>,
                <TableHead key={`${model}-output`}>Output</TableHead>,
              ])}
            </TableRow>
          </TableHeader>
          <TableBody>
            {sortedFlattenedData.map((item) => (
              <TableRow key={`${item.provider}-${item.name}`}>
                <TableCell>
                  {" "}
                  <a href={item.uri} className="underline">
                    {item.provider}
                  </a>
                </TableCell>
                <TableCell>{item.name}</TableCell>

                <TableCell>{item.inputPrice.toFixed(2)}</TableCell>
                <TableCell>{item.outputPrice.toFixed(2)}</TableCell>

                <TableCell className="font-bold">
                  $
                  {(
                    calculatePrice(item.inputPrice, inputTokens) +
                    calculatePrice(item.outputPrice, outputTokens)
                  ).toFixed(2)}
                </TableCell>

                {comparisonModels.flatMap((comparisonModel) => {
                  const [comparisonProvider, comparisonModelName] =
                    comparisonModel.split(":");
                  const comparisonModelData = data
                    .find((p) => p.provider === comparisonProvider)
                    ?.models.find((m) => m.name === comparisonModelName)!;
                  return [
                    <TableCell
                      key={`${comparisonModel}-input`}
                      className={`${
                        parseFloat(
                          calculateComparison(
                            item.inputPrice,
                            comparisonModelData.inputPrice
                          )
                        ) < 0
                          ? "bg-green-100"
                          : parseFloat(
                              calculateComparison(
                                item.inputPrice,
                                comparisonModelData.inputPrice
                              )
                            ) > 0
                          ? "bg-red-100"
                          : ""
                      }`}
                    >
                      {`${item.provider}:${item.name}` === comparisonModel
                        ? "0.00%"
                        : `${calculateComparison(
                            item.inputPrice,
                            comparisonModelData.inputPrice
                          )}%`}
                    </TableCell>,
                    <TableCell
                      key={`${comparisonModel}-output`}
                      className={`${
                        parseFloat(
                          calculateComparison(
                            item.outputPrice,
                            comparisonModelData.outputPrice
                          )
                        ) < 0
                          ? "bg-green-100"
                          : parseFloat(
                              calculateComparison(
                                item.outputPrice,
                                comparisonModelData.outputPrice
                              )
                            ) > 0
                          ? "bg-red-100"
                          : ""
                      }`}
                    >
                      {`${item.provider}:${item.name}` === comparisonModel
                        ? "0.00%"
                        : `${calculateComparison(
                            item.outputPrice,
                            comparisonModelData.outputPrice
                          )}%`}
                    </TableCell>,
                  ];
                })}
              </TableRow>
            ))}
          </TableBody>
        </Table>
      </CardContent>
    </Card>
  );
};

export default App;