frontend/src/components/VideoGallery.tsx

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

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

const VideoGallery: React.FC<GalleryProps> = ({ selectedModel, selectedAttack, examples }) => {
  const exampleItems = examples[selectedModel][selectedAttack]
  const grouped = groupByNameAndVariant(exampleItems)
  const videoNames = Object.keys(grouped)
  const [selectedVideo, setSelectedVideo] = React.useState(videoNames[0] || '')
  const variants = grouped[selectedVideo] || {}
  const variantKeys = Object.keys(variants)
  const [selectedVariant, setSelectedVariant] = React.useState(variantKeys[0] || '')
  const [toggleMode, setToggleMode] = React.useState<'wmd' | 'attacked'>('wmd')
  // Add state for rewind seconds
  const [rewindSeconds, setRewindSeconds] = React.useState(0.5)
  // State for video scale
  const [videoScale, setVideoScale] = React.useState(1)

  // Playback time ref for syncing position
  const playbackTimeRef = React.useRef(0)
  // Refs for all video elements
  const videoRefs = React.useMemo(() => {
    const refs: Record<string, React.RefObject<HTMLVideoElement>> = {}
    variantKeys.forEach((v) => {
      refs[v] = React.createRef<HTMLVideoElement>()
    })
    return refs
  }, [variantKeys.join(',')])

  // Track which variant is currently playing
  const [playingVariant, setPlayingVariant] = React.useState<string | null>(null)

  // Play handler: pause all others, sync time
  const handlePlay = (variant: string) => {
    setPlayingVariant(variant)
    variantKeys.forEach((v) => {
      if (v !== variant && videoRefs[v]?.current) {
        videoRefs[v]?.current?.pause()
      }
    })
  }
  // Pause handler
  const handlePause = (variant: string) => {
    if (playingVariant === variant) setPlayingVariant(null)
  }

  // When selectedVariant changes, sync playback position, rewind, and play state
  React.useEffect(() => {
    if (!selectedVariant) return
    // Rewind playbackTimeRef by rewindSeconds, clamp to 0
    playbackTimeRef.current = Math.max(0, playbackTimeRef.current - rewindSeconds)
    // Set all videos to the new time, pause all except selected
    variantKeys.forEach((v) => {
      const ref = videoRefs[v]?.current
      if (ref) {
        ref.currentTime = playbackTimeRef.current
        if (v !== selectedVariant) {
          ref.pause()
        }
      }
    })
    // If a video was playing, continue playing the swapped variant
    if (playingVariant && videoRefs[selectedVariant]?.current) {
      videoRefs[selectedVariant].current.play()
    }
  }, [selectedVariant])

  // When the selected video plays, update playbackTimeRef
  const handleTimeUpdate = (variantKey: string) => {
    const ref = videoRefs[variantKey]?.current
    if (ref && variantKey === selectedVariant) {
      playbackTimeRef.current = ref.currentTime
    }
  }

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

  if (!videoNames.length) {
    return (
      <div className="w-full mt-12 flex items-center justify-center">
        <div className="text-gray-500">
          No video 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">Video Example</legend>
            <select
              className="select select-bordered"
              value={selectedVideo || ''}
              onChange={(e) => {
                setSelectedVideo(e.target.value || '')
              }}
            >
              {videoNames.map((name) => (
                <option key={name} value={name}>
                  {name}
                </option>
              ))}
            </select>
          </fieldset>
        </div>
        {selectedVideo && 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}
              />
              <div className="flex items-center gap-4 mt-2">
                <label htmlFor="rewind-seconds" className="font-mono text-xs">Rewind Seconds:</label>
                <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 input-xs w-20"
                  placeholder="Seconds"
                />
                <label htmlFor="video-scale" className="font-mono text-xs ml-4">Scale:</label>
                <input
                  id="video-scale"
                  type="range"
                  min={0.3}
                  max={1}
                  step={0.01}
                  value={videoScale}
                  onChange={e => setVideoScale(Number(e.target.value))}
                  className="range range-xs w-40"
                  style={{ verticalAlign: 'middle' }}
                />
                <span className="ml-2 font-mono text-xs">{(videoScale * 100).toFixed(0)}%</span>
              </div>
              <div className="flex flex-col items-center gap-4">
                {variantKeys.map((variantKey) =>
                  variants[variantKey].video_url ? (
                    <video
                      key={variantKey}
                      ref={videoRefs[variantKey]}
                      controls
                      src={variants[variantKey].video_url}
                      className="example-video"
                      style={{
                        width: `${videoScale * 100}%`,
                        height: 'auto',
                        display: selectedVariant === variantKey ? 'block' : 'none',
                        maxWidth: '100%',
                      }}
                      onTimeUpdate={() => handleTimeUpdate(variantKey)}
                      onPlay={() => handlePlay(variantKey)}
                      onPause={() => handlePause(variantKey)}
                    />
                  ) : null
                )}
              </div>
            </ExampleDetailsSection>
          </>
        )}
      </div>
    </div>
  )
}

export default VideoGallery