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 { 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)

  useEffect(() => {
    setData(mockData)
    // Set default comparison models
    setComparisonModels(['GPT-4o', 'Claude 3.5 (Sonnet)', '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)
  }

  return (
    <Card className="w-full max-w-6xl mx-auto">
      <CardHeader>
        <CardTitle>LLM Pricing Comparison Tool</CardTitle>
      </CardHeader>
      <CardContent>
        <div className="">
          <h3 className="text-lg font-semibold mb-2">Select Comparison Models</h3>

          <div className="flex gap-4 mb-4">
            <div className="min-w-96 ">
              <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>

            <div>
              <div className="flex flex-wrap gap-4">
                {data
                  .flatMap((provider) => provider.models)
                  .map((model) => (
                    <div key={model.name} className="flex items-center space-x-2">
                      <Checkbox
                        id={model.name}
                        checked={comparisonModels.includes(model.name)}
                        onCheckedChange={(checked) => {
                          if (checked) {
                            setComparisonModels((prev) => [...prev, model.name])
                          } else {
                            setComparisonModels((prev) => prev.filter((m) => m !== model.name))
                          }
                        }}
                      />
                      <label htmlFor={model.name} className="text-sm font-medium text-gray-700">
                        {model.name}
                      </label>
                    </div>
                  ))}
              </div>
            </div>
          </div>
          <p className="italic text-sm text-muted-foreground">
            Note: If you use Amazon Bedrock or Azure prices for Anthropic, Cohere or OpenAI should be the same.
          </p>
          <Table>
            <TableHeader>
              <TableRow>
                <TableHead rowSpan={2}>Provider</TableHead>
                <TableHead rowSpan={2}>Model</TableHead>
                <TableHead rowSpan={2}>Input Price (per 1M tokens)</TableHead>
                <TableHead rowSpan={2}>Output Price (per 1M tokens)</TableHead>
                <TableHead rowSpan={2}>Total Price</TableHead>
                {comparisonModels.map((model) => (
                  <TableHead key={model} colSpan={2}>
                    Compared to {model}
                  </TableHead>
                ))}
              </TableRow>
              <TableRow>
                {comparisonModels.flatMap((model) => [
                  <TableHead key={`${model}-input`}>Input</TableHead>,
                  <TableHead key={`${model}-output`}>Output</TableHead>,
                ])}
              </TableRow>
            </TableHeader>
            <TableBody>
              {data.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 border border-r-1 border-r-slate-600 border-l-0">
                      $
                      {(
                        calculatePrice(model.inputPrice, inputTokens) + calculatePrice(model.outputPrice, outputTokens)
                      ).toFixed(2)}
                    </TableCell>
                    {comparisonModels.flatMap((comparisonModel) => {
                      const comparisonModelData = data.flatMap((p) => p.models).find((m) => m.name === comparisonModel)!
                      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'
                              : ''
                          }`}
                        >
                          {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'
                              : ''
                          }`}
                        >
                          {model.name === comparisonModel
                            ? '0.00%'
                            : `${calculateComparison(model.outputPrice, comparisonModelData.outputPrice)}%`}
                        </TableCell>,
                      ]
                    })}
                  </TableRow>
                ))
              )}
            </TableBody>
          </Table>
        </div>
      </CardContent>
    </Card>
  )
}

export default App