src/components/BenchmarkTable.tsx

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

interface BenchmarkTableProps {
    data: FlattenedModel[];
    comparisonMetrics?: string[];
    requestSort: (key: string) => void;
    sortConfig: {
        key: string;
        direction: "ascending" | "descending";
    } | null;
    allProviders: Provider[];
    selectedProviders: string[];
    selectedModels: string[];
    onProviderChange: (values: string[]) => void;
    onModelChange: (values: string[]) => void;
    linkProviderModel: boolean;
    onLinkToggle: (checked: boolean) => void;
}

export const BenchmarkTable: React.FC<BenchmarkTableProps> = ({
    data,
    comparisonMetrics,
    requestSort,
    sortConfig,
    allProviders,
    selectedProviders,
    selectedModels,
    onProviderChange,
    onModelChange,
    linkProviderModel,
    onLinkToggle,
}) => {
    const benchmarkMetrics = React.useMemo(() => {
        if (!comparisonMetrics || comparisonMetrics.length === 0) return [];
        return comparisonMetrics;
    }, [comparisonMetrics]);

    const getCellStyle = (value?: number) => {
        if (value === undefined) return "";
        if (value >= 85) return "bg-green-100";
        if (value >= 60) return "bg-yellow-100";
        return "bg-red-100";
    };

    const renderSortIndicator = (key: string) => {
        if (sortConfig?.key !== key) return null;
        return sortConfig.direction === "ascending" ? " ▲" : " ▼";
    };

    const modelOptions = React.useMemo(() => {
        const flat = allProviders.flatMap((provider) =>
            provider.models.map((model) => ({
                label: model.name,
                value: model.name,
                provider: provider.provider,
            }))
        );

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

        return Array.from(new Map(filtered.map((m) => [m.value, m])).values());
    }, [allProviders, selectedProviders, linkProviderModel]);

    return (
        <>
            <div className="flex items-center space-x-4 mb-4">
                <div className="flex-1">
                    <label className="text-sm font-medium block mb-1">Providers</label>
                    <MultiSelect
                        options={allProviders.map((p) => ({ label: p.provider, value: p.provider }))}
                        defaultValue={selectedProviders}
                        onValueChange={onProviderChange}
                    />
                </div>
                <div className="flex-1">
                    <label className="text-sm font-medium block mb-1">Models</label>
                    <MultiSelect
                        options={modelOptions.map(({ label, value }) => ({ label, value }))}
                        defaultValue={selectedModels}
                        onValueChange={onModelChange}
                    />
                </div>
                <div className="flex items-center pt-6 space-x-2">
                    <Switch
                        id="linkBenchmarkProviderModel"
                        checked={linkProviderModel}
                        onCheckedChange={onLinkToggle}
                    />
                    <label htmlFor="linkBenchmarkProviderModel" className="text-sm">
                        Link Provider and Model
                    </label>
                </div>
            </div>

            <Table>
                <TableHeader>
                    <TableRow>
                        <TableHead>
                            <button
                                onClick={() => requestSort("provider")}
                                className="font-medium underline underline-offset-2"
                            >
                                Provider{renderSortIndicator("provider")}
                            </button>
                        </TableHead>
                        <TableHead>
                            <button
                                onClick={() => requestSort("name")}
                                className="font-medium underline underline-offset-2"
                            >
                                Model{renderSortIndicator("name")}
                            </button>
                        </TableHead>
                        <TableHead>
                            <button
                                onClick={() => requestSort("inputPrice")}
                                className="font-medium underline underline-offset-2"
                            >
                                Input Price (USD){renderSortIndicator("inputPrice")}
                            </button>
                        </TableHead>
                        <TableHead>
                            <button
                                onClick={() => requestSort("outputPrice")}
                                className="font-medium underline underline-offset-2"
                            >
                                Output Price (USD){renderSortIndicator("outputPrice")}
                            </button>
                        </TableHead>

                        {benchmarkMetrics.map((metric) => (
                            <TableHead key={metric}>
                                <button
                                    onClick={() => requestSort(metric)}
                                    className="font-medium underline underline-offset-2"
                                >
                                    {metric.replace(/_/g, " ").toUpperCase()}
                                    {renderSortIndicator(metric)}
                                </button>
                            </TableHead>
                        ))}
                    </TableRow>
                </TableHeader>
                <TableBody>
                    {data.map((model) => (
                        <TableRow key={`${model.provider}-${model.name}`}>
                            <TableCell>
                                <a href={model.uri} className="underline">
                                    {model.provider}
                                </a>
                            </TableCell>
                            <TableCell>{model.name}</TableCell>
                            <TableCell>${model.inputPrice.toFixed(2)}</TableCell>
                            <TableCell>${model.outputPrice.toFixed(2)}</TableCell>

                            {benchmarkMetrics.map((metric) => {
                                const score = model.benchmark?.[metric];
                                return (
                                    <TableCell key={metric} className={getCellStyle(score)}>
                                        {score !== undefined ? `${score.toFixed(1)}%` : "-"}
                                    </TableCell>
                                );
                            })}
                        </TableRow>
                    ))}
                </TableBody>
            </Table>
        </>
    );
};