import * as React from "react";
import {
  Table,
  TableHeader,
  TableBody,
  TableRow,
  TableHead,
  TableCell,
} from "@/components/ui/table";
import { MultiSelect } from "@/components/ui/multi-select";
import { FlattenedModel, Provider } from "@/App";

interface PricingTableProps {
  data: FlattenedModel[];
  providers: Provider[];
  selectedProviders: string[];
  selectedModels: string[];
  onProviderChange: (values: string[]) => void;
  onModelChange: (values: string[]) => void;
  comparisonModels: string[];
  inputTokens: number;
  outputTokens: number;
  tokenCalculation: string;
  requestSort: (key: keyof FlattenedModel) => void;
  sortConfig: {
    key: keyof FlattenedModel;
    direction: string;
  } | null;
  linkProviderModel: boolean;
}

export const PricingTable: React.FC<PricingTableProps> = ({
  data,
  providers,
  selectedProviders,
  selectedModels,
  onProviderChange,
  onModelChange,
  comparisonModels,
  inputTokens,
  outputTokens,
  tokenCalculation,
  requestSort,
  sortConfig,
  linkProviderModel,
}) => {
  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 => {
    if (comparisonPrice === 0) return "∞";
    return (((modelPrice - comparisonPrice) / comparisonPrice) * 100).toFixed(2);
  };

  const getModelsForSelectedProviders = () => {
    const flatModels = providers.flatMap((provider) =>
      provider.models.map((model) => ({
        label: model.name,
        value: model.name,
        provider: provider.provider,
      }))
    );

    const filtered = linkProviderModel
      ? flatModels.filter((m) => selectedProviders.includes(m.provider))
      : flatModels;

    // Remove duplicates
    return Array.from(new Map(filtered.map((m) => [m.value, m])).values());
  };

  return (
    <Table>
      <TableHeader>
        <TableRow>
          <TableHead>
            <button onClick={() => requestSort("provider")}>
              Provider {sortConfig?.key === "provider" ? (sortConfig.direction === "ascending" ? "▲" : "▼") : null}
            </button>
          </TableHead>
          <TableHead>
            <button onClick={() => requestSort("name")}>
              Model {sortConfig?.key === "name" ? (sortConfig.direction === "ascending" ? "▲" : "▼") : null}
            </button>
          </TableHead>
          <TableHead>
            <button onClick={() => requestSort("inputPrice")}>
              Input Price (million tokens) {sortConfig?.key === "inputPrice" ? (sortConfig.direction === "ascending" ? "▲" : "▼") : null}
            </button>
          </TableHead>
          <TableHead>
            <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={providers.map((p) => ({ label: p.provider, value: p.provider }))}
              defaultValue={selectedProviders}
              onValueChange={onProviderChange}
            />
          </TableHead>
          <TableHead>
            <MultiSelect
              options={getModelsForSelectedProviders()}
              defaultValue={selectedModels}
              onValueChange={onModelChange}
            />
          </TableHead>
          <TableHead />
          <TableHead />
          <TableHead />
          {comparisonModels.flatMap((model) => [
            <TableHead key={`${model}-input`}>Input</TableHead>,
            <TableHead key={`${model}-output`}>Output</TableHead>,
          ])}
        </TableRow>
      </TableHeader>
      <TableBody>
        {data.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((model) => {
              const [comparisonProvider, comparisonModelName] = model.split(":");
              const comparisonModel = providers
                .find((p) => p.provider === comparisonProvider)
                ?.models.find((m) => m.name === comparisonModelName);

              if (!comparisonModel) return [<TableCell key="missing">–</TableCell>, <TableCell key="missing2">–</TableCell>];

              const inputDelta = calculateComparison(item.inputPrice, comparisonModel.inputPrice);
              const outputDelta = calculateComparison(item.outputPrice, comparisonModel.outputPrice);

              return [
                <TableCell key={`${model}-input`} className={parseFloat(inputDelta) < 0 ? "bg-green-100" : parseFloat(inputDelta) > 0 ? "bg-red-100" : ""}>
                  {`${item.provider}:${item.name}` === model ? "0.00%" : `${inputDelta}%`}
                </TableCell>,
                <TableCell key={`${model}-output`} className={parseFloat(outputDelta) < 0 ? "bg-green-100" : parseFloat(outputDelta) > 0 ? "bg-red-100" : ""}>
                  {`${item.provider}:${item.name}` === model ? "0.00%" : `${outputDelta}%`}
                </TableCell>,
              ];
            })}
          </TableRow>
        ))}
      </TableBody>
    </Table>
  );
};