'use client'; import { useState, useCallback, useEffect, useRef } from 'react'; import { cn } from '@/lib/utils'; import { Button } from '@/components/ui/button'; import { ImageGenCard } from './ImageGenCard'; import { useTranslation } from '@/hooks/useTranslation'; import { usePanel } from '@/hooks/usePanel'; import type { TranslationKey } from '@/i18n'; import type { ReferenceImage } from '@/types'; import type { ImageGenResult } from '@/hooks/useImageGen'; const ASPECT_RATIOS = [ '1:1', '16:9', '9:16', '3:2', '2:3', '4:3', '3:4', '4:5', '5:4', '21:9', ] as const; const RESOLUTIONS = ['1K', '2K', '4K'] as const; interface ImageGenConfirmationProps { messageId?: string; initialPrompt: string; initialAspectRatio: string; initialResolution: string; referenceImages?: ReferenceImage[]; } type Status = 'idle' | 'generating' | 'completed' | 'error'; /** Stable key for persisting generation results in localStorage */ function storageKey(prompt: string, sessionId?: string): string { const prefix = sessionId ? `${sessionId}:` : ''; return `imggen:${prefix}${prompt.slice(0, 80)}`; } export function ImageGenConfirmation({ messageId, initialPrompt, initialAspectRatio, initialResolution, referenceImages, }: ImageGenConfirmationProps) { const { t } = useTranslation(); const { sessionId } = usePanel(); const [prompt, setPrompt] = useState(initialPrompt); const [aspectRatio, setAspectRatio] = useState( ASPECT_RATIOS.includes(initialAspectRatio as typeof ASPECT_RATIOS[number]) ? initialAspectRatio : '1:1' ); const [resolution, setResolution] = useState( RESOLUTIONS.includes(initialResolution as typeof RESOLUTIONS[number]) ? initialResolution : '1K' ); const [status, setStatus] = useState('idle'); const [result, setResult] = useState(null); const [error, setError] = useState(null); const abortRef = useRef(null); // Restore completed result from localStorage on mount useEffect(() => { try { const saved = localStorage.getItem(storageKey(initialPrompt, sessionId)); if (saved) { const parsed: ImageGenResult = JSON.parse(saved); if (parsed.images && parsed.images.length > 0) { setResult(parsed); setStatus('completed'); } } } catch { // ignore } }, [initialPrompt]); const handleStop = useCallback(() => { if (abortRef.current) { abortRef.current.abort(); abortRef.current = null; } setStatus('idle'); }, []); const handleGenerate = useCallback(async () => { const controller = new AbortController(); abortRef.current = controller; setStatus('generating'); setError(null); try { // Split unified ReferenceImage[] back into base64 data vs file paths for the API const refData = referenceImages?.filter(r => r.data).map(r => ({ mimeType: r.mimeType, data: r.data! })); const refPaths = referenceImages?.filter(r => r.localPath).map(r => r.localPath!); const res = await fetch('/api/media/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ prompt, aspectRatio, imageSize: resolution, sessionId, ...(refData && refData.length > 0 ? { referenceImages: refData } : {}), ...(refPaths && refPaths.length > 0 ? { referenceImagePaths: refPaths } : {}), }), signal: controller.signal, }); if (!res.ok) { const err = await res.json().catch(() => ({ error: 'Generation failed' })); throw new Error(err.error || 'Generation failed'); } const data = await res.json(); const genResult: ImageGenResult = { id: data.id, text: data.text, images: data.images || [], }; if (genResult.images.length > 0) { setResult(genResult); setStatus('completed'); // Store lightweight version in localStorage try { const storable = { id: genResult.id, text: genResult.text, images: genResult.images.map(img => ({ mimeType: img.mimeType, localPath: img.localPath, data: '', })), }; localStorage.setItem(storageKey(initialPrompt, sessionId), JSON.stringify(storable)); } catch { // storage full } // Persist result to DB by replacing image-gen-request with image-gen-result { const resultBlock = JSON.stringify({ status: 'completed', prompt, aspectRatio, resolution, images: genResult.images.map(img => ({ mimeType: img.mimeType, localPath: img.localPath, })), }); fetch('/api/chat/messages', { method: 'PUT', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ message_id: messageId || '', content: '```image-gen-result\n' + resultBlock + '\n```', session_id: sessionId, prompt_hint: initialPrompt, }), }).catch(() => {}); } // Defer event dispatch so React commits setResult/setStatus before // ChatView's handler calls sendMessage and triggers a re-render setTimeout(() => { window.dispatchEvent(new CustomEvent('image-gen-completed', { detail: { prompt, aspectRatio, resolution, id: genResult.id, images: genResult.images, }, })); }, 0); } else { setError('No images were generated'); setStatus('error'); } } catch (err) { if ((err as Error).name === 'AbortError') { setStatus('idle'); return; } setError((err as Error).message || 'Generation failed'); setStatus('error'); } finally { abortRef.current = null; } }, [prompt, aspectRatio, resolution, initialPrompt, sessionId, referenceImages]); const handleRegenerate = useCallback(() => { setResult(null); setStatus('idle'); try { localStorage.removeItem(storageKey(initialPrompt, sessionId)); } catch { // ignore } }, [initialPrompt]); // ── Completed: show result only ── if (status === 'completed' && result && result.images.length > 0) { return (
r.data).map(r => ({ mimeType: r.mimeType, data: r.data! }))} />
); } // ── Idle / Generating / Error: show params card ── return (
{/* Header */}
{t('imageGen.confirmTitle' as TranslationKey)}
{/* Reference images preview — unified loop over all reference images */} {referenceImages && referenceImages.length > 0 && (
{referenceImages.map((img, i) => (
{/* eslint-disable-next-line @next/next/no-img-element */} {`Reference
))}
)} {/* Prompt textarea */}