import React, { useState, useEffect, useRef } from 'react'
import { useLocation, useNavigate } from 'react-router-dom'
import API from '../API'
import LoadingSpinner from './LoadingSpinner'
import ImageGallery from './ImageGallery'
import AudioGallery from './AudioGallery'
import VideoGallery from './VideoGallery'
import ModelInfoIcon from './ModelInfoIcon'
import Descriptions from '../Descriptions'
import DatasetSelector from './DatasetSelector'

interface ExamplesProps {
  fileType: 'image' | 'audio' | 'video'
}

// Move ExamplesData type export to allow import in Galleries.tsx
export type ExamplesData = {
  image_url: string
  audio_url?: string
  video_url?: string
  name: string
  metadata: {
    [key: string]: string | boolean
  }
}

const Examples = ({ fileType }: ExamplesProps) => {
  const [examples, setExamples] = useState<{
    [model: string]: { [attack: string]: ExamplesData[] }
  }>({})
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const [selectedModel, setSelectedModel] = useState<string | null>(null)
  const [selectedAttack, setSelectedAttack] = useState<string | null>(null)
  const [descriptionsLoaded, setDescriptionsLoaded] = useState(false)
  const [datasets, setDatasets] = useState<any>({})
  const [selectedDataset, setSelectedDataset] = useState<string>('')
  const descriptions = useRef(Descriptions.getInstance())

  useEffect(() => {
    descriptions.current.load().then(() => setDescriptionsLoaded(true))
    // Fetch datasets when component loads
    API.fetchDatasets().then((datasetsData) => {
      setDatasets(datasetsData)
      // Set default selected dataset based on fileType
      const datasetsForType = datasetsData[fileType] || []
      if (datasetsForType.length > 0) {
        setSelectedDataset(datasetsForType[0])
      }
    }).catch((err) => {
      console.error('Failed to fetch datasets:', err)
    })
  }, [fileType])

  const location = useLocation()
  // Parse query params for model and attack
  useEffect(() => {
    const params = new URLSearchParams(location.search)
    const modelParam = params.get('model')
    const attackParam = params.get('attack')
    const strengthParam = params.get('strength')
    if (modelParam) setSelectedModel(modelParam)
    // Find the most appropriate selectedAttack
    if (attackParam || strengthParam) {
      setSelectedAttack((prev) => {
        if (!selectedModel || !examples[selectedModel]) return prev
        const attacks = Object.keys(examples[selectedModel])
        // If both attack and strength are present, look for attack containing attackParam and ending with strengthParam
        if (attackParam && strengthParam) {
          const found = attacks.find((a) => a.includes(attackParam) && a.endsWith(strengthParam))
          if (found) return found
        }
        // If only attack is present, look for attack containing attackParam
        if (attackParam) {
          const found = attacks.find((a) => a.includes(attackParam))
          if (found) return found
        }
        // If only strength is present, look for attack ending with strengthParam
        if (strengthParam) {
          const found = attacks.find((a) => a.endsWith(strengthParam))
          if (found) return found
        }
        return prev
      })
    }
  }, [location.search, selectedModel, examples])

  useEffect(() => {
    if (!selectedDataset) return

    setLoading(true)
    setError(null)
    API.fetchExamplesByType(fileType, selectedDataset)
      .then((data) => {
        setExamples(data)
        const models = Object.keys(data)
        if (models.length > 0) {
          // If query param exists and is valid, use it, else default to first
          setSelectedModel((prev) => (prev && data[prev] ? prev : models[0]))
          const attacks = Object.keys(data[models[0]])
          if (attacks.length > 0) {
            setSelectedAttack((prev) => (prev && data[models[0]][prev] ? prev : attacks[0]))
          } else {
            setSelectedAttack(null)
          }
        } else {
          setSelectedModel(null)
          setSelectedAttack(null)
        }
        setLoading(false)
      })
      .catch((err) => {
        setError(err.message)
        setLoading(false)
      })
  }, [fileType, selectedDataset])

  if (loading) {
    return <LoadingSpinner />
  }

  return (
    <div className="examples-container">
      <div className="selectors-container flex flex-col gap-4">
        <DatasetSelector
          datasetNames={datasets[fileType] || []}
          selectedDatasetName={selectedDataset}
          onDatasetNameChange={setSelectedDataset}
        />
        <fieldset className="fieldset w-full p-4 rounded border border-gray-700 bg-base-200">
          <legend className="fieldset-legend font-semibold">Model</legend>
          <div className="grid grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-1 max-h-48 overflow-y-auto pr-2">
            {Object.keys(examples).map((model) => {
              const fullName = descriptions.current.getModelFullName(model) || model
              const modelAlias = descriptions.current.getModelAlias(model) || model
              return (
                <div key={model} className="flex items-center gap-2 text-sm relative group">
                  <label className="flex items-center gap-2 flex-grow">
                    <input
                      type="radio"
                      className="radio radio-sm"
                      checked={selectedModel === model}
                      onChange={() => setSelectedModel(model)}
                      name="model-selection"
                    />
                    <div className="flex items-center">
                      <span className="truncate" title={fullName}>
                        {modelAlias}
                      </span>
                      <ModelInfoIcon modelName={model} />
                    </div>
                  </label>
                </div>
              )
            })}
          </div>
        </fieldset>

        {selectedModel && (
          <fieldset className="fieldset">
            <legend className="fieldset-legend">Attack</legend>
            <select
              className="select select-bordered"
              value={selectedAttack || ''}
              onChange={(e) => setSelectedAttack(e.target.value || null)}
            >
              {Object.keys(examples[selectedModel]).map((attack) => (
                <option key={attack} value={attack}>
                  {attack}
                </option>
              ))}
            </select>
          </fieldset>
        )}
      </div>

      {error && <p className="error">Error: {error}</p>}

      {selectedModel && selectedAttack && fileType === 'image' && (
        <ImageGallery
          selectedModel={selectedModel}
          selectedAttack={selectedAttack}
          examples={examples}
        />
      )}
      {selectedModel && selectedAttack && fileType === 'audio' && (
        <AudioGallery
          selectedModel={selectedModel}
          selectedAttack={selectedAttack}
          examples={examples}
        />
      )}
      {selectedModel && selectedAttack && fileType === 'video' && (
        <VideoGallery
          selectedModel={selectedModel}
          selectedAttack={selectedAttack}
          examples={examples}
        />
      )}
    </div>
  )
}

export default Examples