frontend/src/components/ImageGallery.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, { handleVariantToggleClick } from './ExampleVariantToggle'

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

const ImageGallery: React.FC<GalleryProps> = ({ selectedModel, selectedAttack, examples }) => {
  const exampleItems = examples[selectedModel][selectedAttack]

  // Group by image name (name after removing variant prefix), and for each, map variant to metrics
  const grouped = groupByNameAndVariant(exampleItems)

  const imageNames = Object.keys(grouped)
  const [selectedImage, setSelectedImage] = React.useState(imageNames[0] || '')
  const variants = grouped[selectedImage] || {}
  const variantKeys = Object.keys(variants)
  const [selectedVariant, setSelectedVariant] = React.useState(variantKeys[0] || '')
  const [toggleMode, setToggleMode] = React.useState<'wmd' | 'attacked'>('wmd')

  const [zoom, setZoom] = React.useState<{ x: number; y: number } | null>(null)
  const [imgDims, setImgDims] = React.useState<{ width: number; height: number }>({
    width: 0,
    height: 0,
  })
  const imgRef = React.useRef<HTMLImageElement | null>(null)

  // Set a fixed zoom box size for all variants
  const ZOOM_BOX_SIZE = 300
  const [zoomLevel, setZoomLevel] = React.useState(2) // 2x default zoom

  // Calculate the minimum width and height across all variants for the selected image
  const variantImages = variantKeys.map((v) => variants[v]?.image_url).filter(Boolean)
  const [naturalDims, setNaturalDims] = React.useState<{ width: number; height: number } | null>(
    null
  )

  React.useEffect(() => {
    let isMounted = true
    if (variantImages.length > 0) {
      Promise.all(
        variantImages.map(
          (src) =>
            new Promise<{ width: number; height: number }>((resolve) => {
              const img = new window.Image()
              img.onload = () => resolve({ width: img.naturalWidth, height: img.naturalHeight })
              img.src = src!
            })
        )
      ).then((dimsArr) => {
        if (!isMounted) return
        const first = dimsArr[0]
        const allSame = dimsArr.every((d) => d.width === first.width && d.height === first.height)
        if (!allSame) {
          setNaturalDims(null)
        } else {
          setNaturalDims(first)
        }
      })
    }
    return () => {
      isMounted = false
    }
  }, [variantImages.join(','), selectedImage])

  // Place the zoom level slider outside of the zoom hover logic
  const zoomSlider = (
    <div className="flex flex-col items-center mb-4 w-full">
      <label className="font-semibold text-xs mb-1">Zoom Level: {zoomLevel}x</label>
      <input
        type="range"
        min={2}
        max={6}
        step={0.1}
        value={zoomLevel}
        onChange={(e) => setZoomLevel(Number(e.target.value))}
        className="range range-xs w-48"
      />
    </div>
  )

  // If no image is selected, show a message
  if (!variantImages.length) {
    return (
      <div className="w-full mt-12 flex items-center justify-center">
        <div className="text-gray-500">
          No images available. Please select another model and attack.
        </div>
      </div>
    )
  }

  // Ensure the main container allows scrolling and images retain their natural size
  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">Image</legend>
            <select
              className="select select-bordered"
              value={selectedImage || ''}
              onChange={(e) => {
                setSelectedImage(e.target.value || '')
                const newVariants = grouped[e.target.value] || {}
                const newVariantKeys = Object.keys(newVariants)
                setSelectedVariant(newVariantKeys[0] || '')
              }}
            >
              {imageNames.map((name) => (
                <option key={name} value={name}>
                  {name}
                </option>
              ))}
            </select>
          </fieldset>
        </div>
        {selectedImage && 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="radio"
                selectedVariant={selectedVariant}
                setSelectedVariant={setSelectedVariant}
                variantKeys={variantKeys}
              />

              {zoomSlider}
              <div
                style={{
                  width: '100%',
                  display: 'flex',
                  flexDirection: 'row',
                  justifyContent: 'center',
                }}
              >
                <div style={{ position: 'relative', flex: 'none', display: 'inline-block' }}>
                  {/* Original image and overlay */}
                  {naturalDims ? (
                    <img
                      ref={imgRef}
                      src={variants[selectedVariant].image_url as string}
                      alt={selectedImage}
                      width={naturalDims.width}
                      height={naturalDims.height}
                      style={{ display: 'block', cursor: 'zoom-in' }}
                      onMouseMove={(e) => {
                        const rect = (e.target as HTMLImageElement).getBoundingClientRect()
                        const x = e.clientX - rect.left
                        const y = e.clientY - rect.top
                        setZoom({ x, y })
                      }}
                      onMouseLeave={() => setZoom(null)}
                      onClick={() =>
                        handleVariantToggleClick(
                          toggleMode,
                          selectedVariant,
                          setSelectedVariant,
                          variantKeys
                        )
                      }
                    />
                  ) : (
                    <div className="text-red-500 font-bold">
                      Image sizes do not match across variants!
                    </div>
                  )}
                  {/* Zoomed box - now overlays on top of the cursor position */}
                  {zoom &&
                    naturalDims &&
                    (() => {
                      const { width, height } = naturalDims
                      const halfBox = ZOOM_BOX_SIZE / (2 * zoomLevel)
                      let centerX = zoom.x
                      let centerY = zoom.y
                      centerX = Math.max(halfBox, Math.min(width - halfBox, centerX))
                      centerY = Math.max(halfBox, Math.min(height - halfBox, centerY))
                      const bgX = centerX * zoomLevel - ZOOM_BOX_SIZE / 2
                      const bgY = centerY * zoomLevel - ZOOM_BOX_SIZE / 2
                      return (
                        <div
                          style={{
                            position: 'absolute',
                            left: zoom.x - ZOOM_BOX_SIZE / 2, // use raw cursor position for overlay
                            top: zoom.y - ZOOM_BOX_SIZE / 2,
                            width: ZOOM_BOX_SIZE,
                            height: ZOOM_BOX_SIZE,
                            boxSizing: 'border-box',
                            backgroundImage: `url(${variants[selectedVariant]?.image_url})`,
                            backgroundPosition: `-${bgX}px -${bgY}px`,
                            backgroundSize: `${width * zoomLevel}px ${height * zoomLevel}px`,
                            backgroundRepeat: 'no-repeat',
                            zIndex: 10,
                            pointerEvents: 'none',
                            cursor: 'zoom-in',
                          }}
                        />
                      )
                    })()}
                </div>
              </div>
            </ExampleDetailsSection>
          </>
        )}
      </div>
    </div>
  )
}

export default ImageGallery