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'

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[]>([])

  useEffect(() => {
    setData(mockData)
    setComparisonModels(['OpenAI:GPT-4o', 'Anthropic:Claude 3.5 (Sonnet)', 'Google Vertex AI:Gemini 1.5 Pro'])
  }, [])

  const calculatePrice = (price: number, tokens: number): number => {
    return price * tokens
  }

  const calculateComparison = (modelPrice: number, comparisonPrice: number): string => {
    return (((modelPrice - comparisonPrice) / comparisonPrice) * 100).toFixed(2)
  }
  const filteredData = data.filter((provider) => {
    if (selectedProviders.length > 0 && !selectedProviders.includes(provider.provider)) {
      return false
    }
    if (selectedModels.length > 0) {
      return provider.models.some((model) => selectedModels.includes(model.name))
    }
    return true
  })

  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 Comparison Tool</CardTitle>
      </CardHeader>
      <CardContent>
        <div className="mb-4">
          <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 (millions)
            </label>
            <Input
              id="inputTokens"
              type="number"
              value={inputTokens}
              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 (millions)
            </label>
            <Input
              id="outputTokens"
              type="number"
              value={outputTokens}
              onChange={(e) => setOutputTokens(Number(e.target.value))}
              className="mt-1"
            />
          </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>Provider</TableHead>
              <TableHead>Model</TableHead>
              <TableHead>Input Price (per 1M tokens)</TableHead>
              <TableHead>Output Price (per 1M tokens)</TableHead>
              <TableHead>Total Price</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 })) || []
                  }
                  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>
            {filteredData.flatMap((provider) =>
              provider.models.map((model) => (
                <TableRow key={`${provider.provider}-${model.name}`}>
                  <TableCell>
                    {' '}
                    <a href={provider.uri} className="underline">
                      {provider.provider}
                    </a>
                  </TableCell>
                  <TableCell>{model.name}</TableCell>
                  <TableCell>${model.inputPrice.toFixed(2)}</TableCell>
                  <TableCell>${model.outputPrice.toFixed(2)}</TableCell>
                  <TableCell className="font-bold">
                    $
                    {(
                      calculatePrice(model.inputPrice, inputTokens) + calculatePrice(model.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(model.inputPrice, comparisonModelData.inputPrice)) < 0
                            ? 'bg-green-100'
                            : parseFloat(calculateComparison(model.inputPrice, comparisonModelData.inputPrice)) > 0
                            ? 'bg-red-100'
                            : ''
                        }`}
                      >
                        {`${provider.provider}:${model.name}` === comparisonModel
                          ? '0.00%'
                          : `${calculateComparison(model.inputPrice, comparisonModelData.inputPrice)}%`}
                      </TableCell>,
                      <TableCell
                        key={`${comparisonModel}-output`}
                        className={`${
                          parseFloat(calculateComparison(model.outputPrice, comparisonModelData.outputPrice)) < 0
                            ? 'bg-green-100'
                            : parseFloat(calculateComparison(model.outputPrice, comparisonModelData.outputPrice)) > 0
                            ? 'bg-red-100'
                            : ''
                        }`}
                      >
                        {`${provider.provider}:${model.name}` === comparisonModel
                          ? '0.00%'
                          : `${calculateComparison(model.outputPrice, comparisonModelData.outputPrice)}%`}
                      </TableCell>,
                    ]
                  })}
                </TableRow>
              ))
            )}
          </TableBody>
        </Table>
      </CardContent>
    </Card>
  )
}

export default App