frontend/src/components/AudioGallery.tsx

import React from 'react'
import AudioPlayer, { AudioPlayerHandle } from './AudioPlayer'
import type { ExamplesData } from './Examples'
import { groupByNameAndVariant, VARIANT_NAME_MAP } from './galleryUtils'
import ExampleDetailsSection from './ExampleDetailsSection'
import ExampleVariantSelector from './ExampleVariantSelector'
import ExampleVariantMetricsTables from './ExampleVariantMetricsTable'
import ExampleVariantToggle, { handleVariantToggleClick } from './ExampleVariantToggle'
import LoadingSpinner from './LoadingSpinner'
import API from '../API'
import { InformationCircleIcon } from '@heroicons/react/24/outline'
import { Tooltip as ReactTooltip } from 'react-tooltip'
import 'react-tooltip/dist/react-tooltip.css'

interface GalleryProps {
  selectedModel: string
  selectedAttack: string
  examples: {
    [model: string]: {
      [attack: string]: ExamplesData[]
    }
  }
}

const AudioGallery: React.FC<GalleryProps> = ({ selectedModel, selectedAttack, examples }) => {
  const exampleItems = examples[selectedModel][selectedAttack]
  const grouped = groupByNameAndVariant(exampleItems)
  const audioNames = Object.keys(grouped)
  const [selectedAudio, setSelectedAudio] = React.useState(audioNames[0] || '')
  const variants = grouped[selectedAudio] || {}
  const variantKeys = Object.keys(variants)
  const [selectedVariant, setSelectedVariant] = React.useState(variantKeys[0] || '')
  const [toggleMode, setToggleMode] = React.useState<'wmd' | 'attacked'>('wmd')

  // Shared playback state
  const playbackTimeRef = React.useRef(0)
  const [playingVariant, setPlayingVariant] = React.useState<string | null>(null)

  // Refs for all players
  const audioRefs = React.useMemo(() => {
    const refs: Record<string, React.RefObject<AudioPlayerHandle>> = {}
    variantKeys.forEach((v) => {
      refs[v] = React.createRef<AudioPlayerHandle>()
    })
    return refs
  }, [variantKeys.join(',')])

  // Add state for rewind seconds
  const [rewindSeconds, setRewindSeconds] = React.useState(0.5)
  const [imageLoading, setImageLoading] = React.useState(true)

  // Play handler: pause all others, sync time
  const handlePlay = (variant: string) => {
    console.log(`Playing variant: ${variant}`)
    setPlayingVariant(variant)
    variantKeys.forEach((v) => {
      if (v !== variant && audioRefs[v]?.current) {
        audioRefs[v]?.current?.pause()
        // audioRefs[v]?.current?.setTime(playbackTimeRef.current)
      }
    })
  }
  // Pause handler
  const handlePause = (variant: string) => {
    console.log(`Pausing variant: ${variant}`)
    if (playingVariant === variant) setPlayingVariant(null)
  }

  React.useEffect(() => {
    setSelectedVariant(variantKeys[0] || '')
  }, [selectedAudio])

  // Reset image loading state when selected variant changes
  React.useEffect(() => {
    if (variants[selectedVariant]?.image_url) {
      setImageLoading(true)
    }
  }, [selectedVariant, variants[selectedVariant]?.image_url])

  // When selectedVariant changes, play that variant and pause others, syncing position
  React.useEffect(() => {
    if (!selectedVariant) {
      return
    }
    if (playingVariant == null) {
      // On page load don't auto play, only when swapping tracks
      return
    }
    // Rewind playbackTimeRef by rewindSeconds, clamp to 0
    playbackTimeRef.current = Math.max(0, playbackTimeRef.current - rewindSeconds)
    setPlayingVariant(selectedVariant)
    variantKeys.forEach((v) => {
      if (v !== selectedVariant) {
        audioRefs[v]?.current?.pause()
      }
      if (audioRefs[v]?.current) {
        audioRefs[v]?.current?.setTime(playbackTimeRef.current)
      }
    })
  }, [selectedVariant])

  console.log(audioRefs[selectedVariant]?.current?.getCurrentTime())

  if (!audioNames.length) {
    return (
      <div className="w-full mt-12 flex items-center justify-center">
        <div className="text-gray-500">
          No audio examples available. Please select another model and attack.
        </div>
      </div>
    )
  }

  return (
    <div className="w-full overflow-auto" style={{ minHeight: '100vh' }}>
      <div className="example-display">
        <div className="mb-4">
          <fieldset className="fieldset">
            <legend className="fieldset-legend">Audio</legend>
            <select
              className="select select-bordered"
              value={selectedAudio || ''}
              onChange={(e) => {
                setSelectedAudio(e.target.value || '')
              }}
            >
              {audioNames.map((name) => (
                <option key={name} value={name}>
                  {name}
                </option>
              ))}
            </select>
          </fieldset>
        </div>
        {selectedAudio && selectedVariant && variants[selectedVariant] && (
          <>
            <ExampleVariantMetricsTables
              variantMetadatas={Object.fromEntries(
                variantKeys.map((v) => [v, variants[v]?.metadata || {}])
              )}
            />

            <ExampleDetailsSection>
              <ExampleVariantSelector
                variantKeys={variantKeys}
                selectedVariant={selectedVariant}
                setSelectedVariant={setSelectedVariant}
              />
              <ReactTooltip
                id="variant-selector-tooltip"
                place="top"
                className="z-[10000] max-w-xs !opacity-100 bg-base-100 text-base-content"
                style={{ boxShadow: '0 0 10px rgba(0,0,0,0.2)', zIndex: 10000 }}
                positionStrategy="fixed"
              >
                <div className="p-2 text-xs text-left relative z-[10000]">
                  You can also change the variant using keys <b>1</b>, <b>2</b>, <b>3</b>, <b>4</b>{' '}
                  on your keyboard.
                </div>
              </ReactTooltip>
              <ExampleVariantToggle
                toggleMode={toggleMode}
                setToggleMode={setToggleMode}
                selectedVariant={selectedVariant}
                setSelectedVariant={setSelectedVariant}
                variantKeys={variantKeys}
              />
              <fieldset className="fieldset mt-2">
                <legend className="fieldset-legend">Rewind Seconds</legend>
                <input
                  id="rewind-seconds"
                  type="number"
                  min={0}
                  step={0.1}
                  value={rewindSeconds}
                  onChange={(e) => setRewindSeconds(Math.max(0, Number(e.target.value)))}
                  className="input input-bordered w-20"
                  placeholder="Seconds"
                />
              </fieldset>
              <div className="flex flex-col items-center gap-4">
                {variantKeys.map((variantKey) =>
                  variants[variantKey].audio_url ? (
                    <div
                      key={variantKey}
                      style={{
                        width: '100%',
                        display: selectedVariant === variantKey ? 'block' : 'none',
                      }}
                    >
                      <div className="font-mono text-xs mb-1">{variantKey}</div>
                      <AudioPlayer
                        ref={audioRefs[variantKey]}
                        src={API.getProxiedUrl(variants[variantKey].audio_url, true)}
                        playing={playingVariant === variantKey}
                        onPlay={() => handlePlay(variantKey)}
                        onPause={() => handlePause(variantKey)}
                        onAudioProcess={(currentTime) => {
                          playbackTimeRef.current = currentTime
                          variantKeys.forEach((v) => {
                            if (v !== variantKey && audioRefs[v]?.current) {
                              audioRefs[v]?.current?.setTime(currentTime)
                            }
                          })
                        }}
                      />
                    </div>
                  ) : null
                )}
                {variants[selectedVariant].image_url && (
                  <div className="relative">
                    {imageLoading && <LoadingSpinner />}
                    <img
                      key={variants[selectedVariant].image_url}
                      src={API.getProxiedUrl(variants[selectedVariant].image_url)}
                      alt={selectedAudio}
                      className="example-image"
                      style={{
                        display: 'block',
                        opacity: imageLoading ? 0 : 1,
                        transition: 'opacity 0.3s',
                      }}
                      onLoad={() => setImageLoading(false)}
                      onError={() => setImageLoading(false)}
                      onClick={() =>
                        handleVariantToggleClick(
                          toggleMode,
                          selectedVariant,
                          setSelectedVariant,
                          variantKeys
                        )
                      }
                    />
                  </div>
                )}
              </div>
            </ExampleDetailsSection>
          </>
        )}
      </div>
    </div>
  )
}

export default AudioGallery