frontend/src/components/ImageGallery.tsx

import React from 'react'
import type { ExamplesData } from './Examples'
import { groupByNameAndVariant, VARIANT_NAME_MAP } from './galleryUtils'
import ExampleVariantMetricsTables from './ExampleVariantMetricsTable'
import ExampleDetailsSection from './ExampleDetailsSection'
import ExampleVariantSelector from './ExampleVariantSelector'
import ExampleVariantToggle, { handleVariantToggleClick } from './ExampleVariantToggle'
import API from '../API'
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 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 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) => API.getProxiedUrl(variants[v]?.image_url))
    .filter(Boolean)
  const [originalImageDims, setOriginalImageDims] = React.useState<{
    width: number
    height: number
  }>({ width: 0, height: 0 })
  const [dimsMismatch, setDimsMismatch] = React.useState(false)

  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)
        setDimsMismatch(!allSame)
        setOriginalImageDims(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={1}
        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] && (
          <>
            <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}
              />

              {zoomSlider}
              <div
                style={{
                  width: '100%',
                  display: 'flex',
                  flexDirection: 'row',
                  justifyContent: 'center',
                }}
              >
                <div style={{ position: 'relative', flex: 'none', display: 'inline-block' }}>
                  {/* Original image and overlay */}
                  {dimsMismatch ? (
                    <div className="text-red-500 font-bold text-center">
                      Image sizes do not match across variants
                    </div>
                  ) : null}

                  <img
                    ref={imgRef}
                    src={API.getProxiedUrl(variants[selectedVariant].image_url as string)}
                    alt={selectedImage}
                    width={originalImageDims.width}
                    height={originalImageDims.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
                      )
                    }
                  />
                  {/* Zoomed box - now overlays on top of the cursor position */}
                  {zoom &&
                    (() => {
                      const { width, height } = originalImageDims
                      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(${API.getProxiedUrl(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