frontend/src/components/AudioGallery.tsx

import React from 'react'
import AudioPlayer, { AudioPlayerHandle } from './AudioPlayer'
import type { ExamplesData } from './Examples'
import { groupByNameAndVariant } from './galleryUtils'
import ExampleDetailsSection from './ExampleDetailsSection'
import ExampleVariantSelector from './ExampleVariantSelector'
import ExampleVariantMetricsTable from './ExampleVariantMetricsTable'
import ExampleVariantToggle, { handleVariantToggleClick } from './ExampleVariantToggle'

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)

  // 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])

  // 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] && (
          <>
            <ExampleVariantMetricsTable
              variantMetadatas={Object.fromEntries(
                variantKeys.map((v) => [v, variants[v]?.metadata || {}])
              )}
            />

            <ExampleDetailsSection>
              <ExampleVariantSelector
                variantKeys={variantKeys}
                selectedVariant={selectedVariant}
                setSelectedVariant={setSelectedVariant}
              />
              <ExampleVariantToggle
                toggleMode={toggleMode}
                setToggleMode={setToggleMode}
                type="button"
                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={variants[variantKey].audio_url}
                        playing={playingVariant === variantKey}
                        onPlay={() => handlePlay(variantKey)}
                        onPause={() => handlePause(variantKey)}
                        onAudioProcess={(currentTime) => {
                          // console.log(`Current time for ${variantKey}: ${currentTime}`)
                          playbackTimeRef.current = currentTime
                          variantKeys.forEach((v) => {
                            if (v !== variantKey && audioRefs[v]?.current) {
                              audioRefs[v]?.current?.setTime(currentTime)
                            }
                          })
                        }}
                      />
                    </div>
                  ) : null
                )}
                {variants[selectedVariant].image_url && (
                  <img
                    src={variants[selectedVariant].image_url}
                    alt={selectedAudio}
                    className="example-image"
                    style={{ display: 'block' }}
                    onClick={() =>
                      handleVariantToggleClick(
                        toggleMode,
                        selectedVariant,
                        setSelectedVariant,
                        variantKeys
                      )
                    }
                  />
                )}
              </div>
            </ExampleDetailsSection>
          </>
        )}
      </div>
    </div>
  )
}

export default AudioGallery