import React, { useState, useEffect } from 'react'
import DatasetSelector from './DatasetSelector'
import LeaderboardTable from './LeaderboardTable'
import LeaderboardChart from './LeaderboardChart'
import ModelFilter from './ModelFilter'
import API from '../API'
import LoadingSpinner from './LoadingSpinner'

interface LeaderboardPageProps {
  datasetType: 'audio' | 'image' | 'video'
}

const DATASET_NAMES: { [type: string]: string[] } = {
  audio: ['voxpopuli_1k/audio', 'ravdess_1k/audio'],
  image: ['val2014_1k/image'],
  video: ['sav_val_full/video'],
}

const LeaderboardPage: React.FC<LeaderboardPageProps> = ({ datasetType }) => {
  const datasetNames = DATASET_NAMES[datasetType]
  const [selectedDatasetName, setSelectedDatasetName] = useState(datasetNames[0])

  // Update selectedDatasetName when datasetType changes
  useEffect(() => {
    setSelectedDatasetName(DATASET_NAMES[datasetType][0])
  }, [datasetType])
  const [models, setModels] = useState<string[]>([])
  const [selectedModels, setSelectedModels] = useState<Set<string>>(new Set())
  const [loading, setLoading] = useState(true)
  const [benchmarkData, setBenchmarkData] = useState<any>(null)

  // Fetch available models when dataset changes
  useEffect(() => {
    setLoading(true)
    API.fetchStaticFile(`data/${selectedDatasetName}?dataset_type=benchmark`)
      .then((response) => {
        const data = JSON.parse(response)
        setBenchmarkData(data)
        const rows = data['rows']
        const allKeys: string[] = Array.from(new Set(rows.flatMap((row: any) => Object.keys(row))))
        // Remove 'metric' from headers if it exists
        const headers = allKeys.filter((key) => key !== 'metric')

        setModels(headers)
        // Initialize all models as selected
        setSelectedModels(new Set(headers))
        setLoading(false)
      })
      .catch((err) => {
        console.error('Failed to fetch models:', err)
        setLoading(false)
      })
  }, [selectedDatasetName])

  return (
    <div className="">
      <div className="flex flex-col gap-4">
        <DatasetSelector
          datasetNames={datasetNames}
          selectedDatasetName={selectedDatasetName}
          onDatasetNameChange={setSelectedDatasetName}
        />
      </div>
      {loading ? (
        <LoadingSpinner />
      ) : (
        <>
          {models.length > 0 && (
            <ModelFilter
              models={models}
              selectedModels={selectedModels}
              setSelectedModels={setSelectedModels}
            />
          )}
          <div className="space-y-8">
            <LeaderboardTable benchmarkData={benchmarkData} selectedModels={selectedModels} />
            <div className="mt-8 pt-4 border-t border-gray-200">
              <LeaderboardChart dataset={selectedDatasetName} selectedModels={selectedModels} />
            </div>
          </div>
        </>
      )}
    </div>
  )
}

export default LeaderboardPage