frontend/src/components/ImageGallery.tsx

import React from 'react'
import type { ExamplesData } from './Examples'
import { groupByNameAndVariant } from './galleryUtils'

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

const Metadata: React.FC<{ metadata: Record<string, string | number | boolean> }> = ({
  metadata,
}) => (
  <fieldset className="fieldset w-full p-4 rounded border border-gray-700 bg-base-200 mb-4">
    <legend className="fieldset-legend font-semibold">Metadata</legend>
    <ul className="text-xs flex flex-row gap-4">
      {(() => {
        const entries = Object.entries(metadata || {}).filter(([k]) => k !== 'image_url')
        const detectedIdx = entries.findIndex(([k]) => k === 'detected')
        if (detectedIdx > -1) {
          const [detected] = entries.splice(detectedIdx, 1)
          entries.unshift(detected)
        }
        return entries.map(([k, v]) => (
          <li key={k}>
            <span className="font-mono">{k}</span>: {String(v)}
          </li>
        ))
      })()}
    </ul>
  </fieldset>
)

const VariantSelector: React.FC<{
  variantKeys: string[]
  selectedVariant: string
  setSelectedVariant: (v: string) => void
}> = ({ variantKeys, selectedVariant, setSelectedVariant }) => (
  <fieldset className="fieldset w-full p-4 rounded border border-gray-700 bg-base-200 mb-4">
    <legend className="fieldset-legend font-semibold">Variants</legend>
    <div className="mb-2 flex gap-4 flex-wrap">
      {variantKeys.map((variant, idx) => (
        <label key={variant} className="flex items-center gap-1 cursor-pointer">
          <input
            type="radio"
            name={`variant-selector`}
            value={variant}
            checked={selectedVariant === variant}
            onChange={() => setSelectedVariant(variant)}
          />
          <span className="text-xs font-semibold">
            {variant} <span className="opacity-60">[{idx + 1}]</span>
          </span>
        </label>
      ))}
    </div>
  </fieldset>
)

const ExampleDetailsSection: React.FC<{ children: React.ReactNode }> = ({ children }) => (
  <fieldset className="fieldset w-full p-4 rounded border border-gray-700 bg-base-200 mt-6">
    {children}
  </fieldset>
)

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 [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

  // Keyboard shortcut for variant switching
  React.useEffect(() => {
    const handler = (e: KeyboardEvent) => {
      if (document.activeElement && (document.activeElement as HTMLElement).tagName === 'INPUT')
        return
      const idx = parseInt(e.key, 10)
      if (!isNaN(idx) && idx > 0 && idx <= variantKeys.length) {
        setSelectedVariant(variantKeys[idx - 1])
      }
    }
    window.addEventListener('keydown', handler)
    return () => window.removeEventListener('keydown', handler)
  }, [variantKeys])

  // 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] && (
          <>
            <VariantSelector
              variantKeys={variantKeys}
              selectedVariant={selectedVariant}
              setSelectedVariant={setSelectedVariant}
            />
            <Metadata metadata={variants[selectedVariant].metadata || {}} />
            {zoomSlider}
            <ExampleDetailsSection>
              <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)}
                    />
                  ) : (
                    <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