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

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 [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 = data
    .filter(
      (provider) =>
        selectedProviders.length === 0 ||
        selectedProviders.includes(provider.provider)
    )
    .map((provider) => ({
      ...provider,
      models: provider.models.filter(
        (model) =>
          selectedModels.length === 0 || 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]
    );
  };

  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>

        <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={
                    data
                      .flatMap((provider) => provider.models)
                      .map((model) => ({ label: model.name, value: model.name }))
                      .reduce(
                        (
                          acc: { label: string; value: string }[],
                          curr: { label: string; value: string }
                        ) => {
                          if (!acc.find((m) => m.value === curr.value)) {
                            acc.push(curr);
                          }
                          return acc;
                        },
                        []
                      ) || []
                  }
                  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;