frontend/src/Variations.jsx

import { useEffect, useState } from 'react'
import JSZip from 'jszip'
import ImageLightbox from './ImageLightbox'

const API_BASE = '/api'
const AUTH_TOKEN_KEY = 'amalfa_auth_token'

function authHeaders() {
  const token = typeof localStorage !== 'undefined' ? localStorage.getItem(AUTH_TOKEN_KEY) : null
  if (!token) return {}
  return { Authorization: `Bearer ${token}` }
}

function proxyImageUrl(url) {
  return url ? `${API_BASE}/proxy-image?url=${encodeURIComponent(url)}` : ''
}

function downloadImage(url, filename) {
  const a = document.createElement('a')
  a.href = proxyImageUrl(url) || url
  a.download = filename
  document.body.appendChild(a)
  a.click()
  a.remove()
}

async function fetchImageBlob(url) {
  if (!url) throw new Error('Missing image URL')
  try {
    const proxied = await fetch(proxyImageUrl(url))
    if (proxied.ok) return await proxied.blob()
  } catch (_) {}
  const direct = await fetch(url)
  if (!direct.ok) throw new Error('Failed to fetch image')
  return await direct.blob()
}

export default function Variations({ onLogout }) {
  const [imageModels, setImageModels] = useState([])
  const [selectedImageModel, setSelectedImageModel] = useState('nano-banana-2')
  const [aspectRatio, setAspectRatio] = useState('1:1')
  const [variationCount, setVariationCount] = useState(10)
  const [userPrompt, setUserPrompt] = useState('')
  const [winningImageUrl, setWinningImageUrl] = useState('')
  const [winningPreviewUrl, setWinningPreviewUrl] = useState('')
  const [uploading, setUploading] = useState(false)
  const [generating, setGenerating] = useState(false)
  const [progress, setProgress] = useState({ done: 0, total: 0 })
  const [results, setResults] = useState([])
  const [selected, setSelected] = useState([])
  const [downloading, setDownloading] = useState(false)
  const [error, setError] = useState(null)
  const [lightboxImage, setLightboxImage] = useState(null)

  useEffect(() => {
    fetch(`${API_BASE}/image-models`, { headers: authHeaders() })
      .then((r) => {
        if (r.status === 401) {
          localStorage.removeItem(AUTH_TOKEN_KEY)
          onLogout?.()
          return null
        }
        return r.ok ? r.json() : null
      })
      .then((data) => {
        if (!data) return
        setImageModels(data.models || [])
        if (data.default) setSelectedImageModel(data.default)
      })
      .catch(() => {})
  }, [onLogout])

  useEffect(() => {
    return () => {
      if (winningPreviewUrl) URL.revokeObjectURL(winningPreviewUrl)
    }
  }, [winningPreviewUrl])

  async function handleUpload(file) {
    if (!file) return
    if (winningPreviewUrl) URL.revokeObjectURL(winningPreviewUrl)
    setWinningPreviewUrl(URL.createObjectURL(file))
    setUploading(true)
    setError(null)
    try {
      const form = new FormData()
      form.append('file', file)
      const res = await fetch(`${API_BASE}/upload-reference`, {
        method: 'POST',
        headers: authHeaders(),
        body: form,
      })
      if (res.status === 401) {
        localStorage.removeItem(AUTH_TOKEN_KEY)
        onLogout?.()
        return
      }
      const data = await res.json().catch(() => ({}))
      if (!res.ok) throw new Error(data.detail || 'Upload failed')
      setWinningImageUrl(data.url || '')
    } catch (e) {
      setError(e.message || 'Upload failed')
    } finally {
      setUploading(false)
    }
  }

  async function handleGenerate() {
    if (!winningImageUrl) {
      setError('Upload a winning creative first.')
      return
    }
    setGenerating(true)
    setError(null)
    setResults([])
    setSelected([])
    setProgress({ done: 0, total: variationCount })
    try {
      const res = await fetch(`${API_BASE}/generate-variations/stream`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', ...authHeaders() },
        body: JSON.stringify({
          image_url: winningImageUrl,
          n: variationCount,
          model_key: selectedImageModel,
          aspect_ratio: aspectRatio,
          user_prompt: userPrompt || undefined,
        }),
      })
      if (res.status === 401) {
        localStorage.removeItem(AUTH_TOKEN_KEY)
        onLogout?.()
        return
      }
      if (!res.ok) {
        const data = await res.json().catch(() => ({}))
        throw new Error(data.detail || 'Variation generation failed')
      }
      const reader = res.body.getReader()
      const decoder = new TextDecoder()
      let buffer = ''
      while (true) {
        const { done, value } = await reader.read()
        if (done) break
        buffer += decoder.decode(value, { stream: true })
        const chunks = buffer.split('\n\n')
        buffer = chunks.pop() || ''
        for (const chunk of chunks) {
          const dataLine = chunk.split('\n').find((l) => l.startsWith('data: '))
          if (!dataLine) continue
          try {
            const payload = JSON.parse(dataLine.slice(6))
            if (payload.event === 'result' && payload.data) {
              const rowId = `${payload.data.variation_id ?? 'v'}-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`
              setResults((prev) => [...prev, { ...payload.data, rowId }])
              setProgress((prev) => ({ ...prev, done: prev.done + 1 }))
            } else if (payload.event === 'error') {
              setError(payload.message || 'Variation generation failed')
            }
          } catch (_) {}
        }
      }
    } catch (e) {
      setError(e.message || 'Variation generation failed')
    } finally {
      setGenerating(false)
    }
  }

  function toggleSelect(rowId) {
    setSelected((prev) => (prev.includes(rowId) ? prev.filter((id) => id !== rowId) : [...prev, rowId]))
  }

  async function handleDownloadSelected() {
    const toDownload = results.filter((r) => r.image_url && selected.includes(r.rowId))
    if (!toDownload.length) return
    setDownloading(true)
    try {
      const zip = new JSZip()
      await Promise.all(
        toDownload.map(async (item, idx) => {
          const blob = await fetchImageBlob(item.image_url)
          const id = String(item.variation_id ?? idx + 1).padStart(2, '0')
          const slug = (item.concept_name || 'variation')
            .replace(/[^a-z0-9]+/gi, '-')
            .toLowerCase()
            .replace(/^-|-$/g, '') || 'variation'
          zip.file(`Amalfa-variation-${id}-${slug}.png`, blob)
        })
      )
      const blob = await zip.generateAsync({ type: 'blob' })
      const a = document.createElement('a')
      a.href = URL.createObjectURL(blob)
      a.download = 'Amalfa-variations-selected.zip'
      a.click()
      URL.revokeObjectURL(a.href)
    } catch (e) {
      setError(e.message || 'Failed to prepare ZIP')
    } finally {
      setDownloading(false)
    }
  }

  return (
    <section className="card variations-card">
      {lightboxImage && (
        <ImageLightbox
          src={lightboxImage.src}
          alt={lightboxImage.alt}
          onClose={() => setLightboxImage(null)}
        />
      )}
      <h2>Winning Creative Variations</h2>
      <p className="gallery-desc">Upload your best-performing creative and generate multiple fresh variations.</p>
      {error && <p className="error">{error}</p>}

      <div className="variations-controls">
        <label>
          <span>Winning creative image</span>
          <input
            type="file"
            accept="image/png,image/jpeg,image/webp"
            onChange={(e) => handleUpload(e.target.files?.[0])}
            disabled={uploading || generating}
          />
        </label>
        <label>
          <span>User prompt (optional)</span>
          <input
            type="text"
            value={userPrompt}
            onChange={(e) => setUserPrompt(e.target.value)}
            placeholder="e.g. Make it moodier, cinematic lighting, more close-up angles"
            disabled={generating}
          />
        </label>
        <label>
          <span>How many variations</span>
          <input
            type="range"
            min="1"
            max="30"
            value={variationCount}
            onChange={(e) => setVariationCount(Number(e.target.value))}
            disabled={generating}
          />
          <strong>{variationCount}</strong>
        </label>
        <label>
          <span>Image model</span>
          <select
            value={selectedImageModel}
            onChange={(e) => setSelectedImageModel(e.target.value)}
            disabled={generating}
          >
            {imageModels.map((m) => (
              <option key={m.key} value={m.key}>
                {m.label || m.key}
              </option>
            ))}
          </select>
        </label>
        <label>
          <span>Aspect ratio</span>
          <select value={aspectRatio} onChange={(e) => setAspectRatio(e.target.value)} disabled={generating}>
            <option value="1:1">1:1</option>
            <option value="16:9">16:9</option>
            <option value="9:16">9:16</option>
          </select>
        </label>
      </div>

      {winningImageUrl && (
        <div className="variations-winning-preview">
          <img src={winningPreviewUrl || winningImageUrl} alt="Winning creative" />
        </div>
      )}

      <div className="variations-actions">
        <button type="button" className="btn-run" disabled={generating || uploading || !winningImageUrl} onClick={handleGenerate}>
          {generating ? 'Generating…' : `Generate ${variationCount} Variations`}
        </button>
        {generating && <span className="gallery-item-date">Progress: {progress.done}/{progress.total}</span>}
      </div>

      {results.length > 0 && (
        <>
          <div className="gallery-toolbar">
            <label className="gallery-select-all">
              <input
                type="checkbox"
                checked={selected.length === results.length && results.length > 0}
                onChange={(e) => (e.target.checked ? setSelected(results.map((r) => r.rowId)) : setSelected([]))}
              />
              Select all
            </label>
            <button
              type="button"
              className="btn-download-zip"
              disabled={downloading || selected.length === 0}
              onClick={handleDownloadSelected}
            >
              {downloading ? 'Preparing…' : `Download selected (${selected.length})`}
            </button>
          </div>
          <div className="gallery-grid">
          {results.map((item, idx) => (
            <div className="gallery-item" key={item.rowId || `${item.variation_id || idx}-${idx}`}>
              <div className="gallery-item-img-wrap">
                {item.image_url ? (
                  <img
                    src={proxyImageUrl(item.image_url)}
                    alt={item.concept_name || `Variation ${idx + 1}`}
                    className="gallery-item-img img-expandable"
                    onClick={() =>
                      setLightboxImage({
                        src: proxyImageUrl(item.image_url),
                        alt: item.concept_name || `Variation ${idx + 1}`,
                      })
                    }
                  />
                ) : (
                  <div className="gallery-item-placeholder">Failed</div>
                )}
                <label className="gallery-item-checkbox">
                  <input
                    type="checkbox"
                    checked={selected.includes(item.rowId)}
                    onChange={() => toggleSelect(item.rowId)}
                  />
                  <span>Select</span>
                </label>
              </div>
              <div className="gallery-item-footer">
                <p className="gallery-item-concept">#{item.variation_id || idx + 1} {item.concept_name || 'Variation'}</p>
                {item.error ? (
                  <p className="error">{item.error}</p>
                ) : (
                  <button
                    type="button"
                    className="btn-gallery-download-one"
                    onClick={() => downloadImage(item.image_url, `variation-${item.variation_id || idx + 1}.png`)}
                  >
                    Download
                  </button>
                )}
              </div>
            </div>
          ))}
        </div>
        </>
      )}
    </section>
  )
}