persona-community-3/apps/creator-ui/src/pages/GeneratePage.tsx

import { useState, useMemo } from 'react';
import { useNavigate } from 'react-router-dom';
import { useAuth } from '@persona-community-3/auth';
import { useMediaGeneration } from '@persona-community-3/realtime';
import {
  Card,
  CardHeader,
  CardTitle,
  CardDescription,
  CardContent,
  Button,
  FormField,
  Select,
  SelectTrigger,
  SelectValue,
  SelectContent,
  SelectItem,
  ImageGrid,
  VideoGrid,
  GenerationProgress,
  ProviderBadge,
  Loader2,
} from '@persona-community-3/ui';

type GenerateMode = 'image' | 'video';

interface ImageResult {
  images: Array<{ data: string; isUrl: boolean; seed?: number }>;
  provider: string;
  latencyMs: number;
}

interface VideoResult {
  videos: Array<{ data: string; isUrl: boolean; mimeType: string }>;
  provider: string;
  latencyMs: number;
}

export function GeneratePage() {
  const { user, getToken } = useAuth();
  const navigate = useNavigate();
  const [mode, setMode] = useState<GenerateMode>('image');
  const [prompt, setPrompt] = useState('');
  const [aspectRatio, setAspectRatio] = useState('1:1');
  const [count, setCount] = useState(1);
  const [duration, setDuration] = useState('5s');

  const apiPrefix = import.meta.env.VITE_API_URL || '';

  const authHeaders = useMemo(() => {
    const token = getToken();
    return token ? { Authorization: `Bearer ${token}` } : undefined;
  }, [getToken]);

  const imageGen = useMediaGeneration<ImageResult>({
    endpoint: `${apiPrefix}/api/{{SERVICE_NAME}}/generate/image`,
    sseEndpoint: `${apiPrefix}/api/{{SERVICE_NAME}}/events`,
    userId: user?.id || 'anonymous',
    headers: authHeaders,
  });

  const videoGen = useMediaGeneration<VideoResult>({
    endpoint: `${apiPrefix}/api/{{SERVICE_NAME}}/generate/video`,
    sseEndpoint: `${apiPrefix}/api/{{SERVICE_NAME}}/events`,
    userId: user?.id || 'anonymous',
    headers: authHeaders,
  });

  const gen = mode === 'image' ? imageGen : videoGen;
  const isGenerating = gen.status === 'pending' || gen.status === 'generating';

  const handleGenerate = async () => {
    if (!prompt.trim()) return;
    gen.reset();
    const request = mode === 'image'
      ? { prompt, count, aspectRatio }
      : { prompt, aspectRatio, duration };
    await gen.generate(request);
  };

  return (
    <div className="space-y-6">
      <Card>
        <CardHeader>
          <CardTitle>AI Generation</CardTitle>
          <CardDescription>
            Generate images and videos using AI (Gemini / LaoZhang)
          </CardDescription>
        </CardHeader>
        <CardContent className="space-y-4">
          {/* Mode toggle */}
          <div className="flex gap-2">
            <Button
              variant={mode === 'image' ? 'default' : 'outline'}
              onClick={() => setMode('image')}
              disabled={isGenerating}
            >
              Images
            </Button>
            <Button
              variant={mode === 'video' ? 'default' : 'outline'}
              onClick={() => {
                setMode('video');
                if (aspectRatio === '1:1') setAspectRatio('16:9');
              }}
              disabled={isGenerating}
            >
              Video
            </Button>
          </div>

          <FormField
            label="Prompt"
            name="prompt"
            value={prompt}
            onChange={(e) => setPrompt(e.target.value)}
            placeholder={
              mode === 'image'
                ? 'A serene mountain landscape at sunset...'
                : 'A cat playing piano in a jazz club...'
            }
          />

          <div className="grid grid-cols-2 gap-4">
            <div className="space-y-2">
              <label className="text-sm font-medium text-[var(--text-primary)]">Aspect Ratio</label>
              <Select value={aspectRatio} onValueChange={(v) => setAspectRatio(v)}>
                <SelectTrigger>
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  {mode === 'image' && <SelectItem value="1:1">Square (1:1)</SelectItem>}
                  <SelectItem value="16:9">Landscape (16:9)</SelectItem>
                  <SelectItem value="9:16">Portrait (9:16)</SelectItem>
                </SelectContent>
              </Select>
            </div>

            {mode === 'image' ? (
              <div className="space-y-2">
                <label className="text-sm font-medium text-[var(--text-primary)]">Count</label>
                <Select value={String(count)} onValueChange={(v) => setCount(Number(v))}>
                  <SelectTrigger>
                    <SelectValue />
                  </SelectTrigger>
                  <SelectContent>
                    <SelectItem value="1">1 image</SelectItem>
                    <SelectItem value="2">2 images</SelectItem>
                    <SelectItem value="4">4 images</SelectItem>
                  </SelectContent>
                </Select>
              </div>
            ) : (
              <div className="space-y-2">
                <label className="text-sm font-medium text-[var(--text-primary)]">Duration</label>
                <Select value={duration} onValueChange={(v) => setDuration(v)}>
                  <SelectTrigger>
                    <SelectValue />
                  </SelectTrigger>
                  <SelectContent>
                    <SelectItem value="5s">5 seconds</SelectItem>
                    <SelectItem value="10s">10 seconds</SelectItem>
                  </SelectContent>
                </Select>
              </div>
            )}
          </div>

          <Button onClick={handleGenerate} disabled={isGenerating || !prompt.trim()}>
            {isGenerating && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
            {isGenerating ? 'Generating...' : `Generate ${mode === 'image' ? 'Images' : 'Video'}`}
          </Button>
        </CardContent>
      </Card>

      {isGenerating && (
        <Card>
          <CardContent className="py-4 space-y-3">
            <div className="flex items-center justify-between text-sm">
              <span className="text-[var(--text-muted)]">{gen.message || 'Starting...'}</span>
              <span className="text-[var(--text-muted)]">{gen.progress}%</span>
            </div>
            <GenerationProgress percent={gen.progress} />
            {gen.sseState !== 'connected' && (
              <p className="text-xs text-[var(--warning)]">
                SSE {gen.sseState} — events may be delayed
              </p>
            )}
          </CardContent>
        </Card>
      )}

      {gen.status === 'failed' && gen.error && (
        <Card className="border-[var(--error)]">
          <CardContent className="py-4 text-[var(--error)]">
            {gen.error}
          </CardContent>
        </Card>
      )}

      {gen.status === 'complete' && imageGen.result && mode === 'image' && (
        <Card>
          <CardHeader className="flex flex-row items-center justify-between">
            <CardTitle>Results</CardTitle>
            <div className="flex items-center gap-2">
              {imageGen.result.provider && <ProviderBadge provider={imageGen.result.provider} />}
              <Button variant="outline" size="sm" onClick={() => navigate('/media')}>
                View in Library
              </Button>
            </div>
          </CardHeader>
          <CardContent>
            <ImageGrid
              images={imageGen.result.images.map((img) => ({
                src: img.isUrl ? img.data : `data:image/png;base64,${img.data}`,
                alt: prompt,
              }))}
              columns={imageGen.result.images.length > 1 ? 2 : 1}
            />
          </CardContent>
        </Card>
      )}

      {gen.status === 'complete' && videoGen.result && mode === 'video' && (
        <Card>
          <CardHeader className="flex flex-row items-center justify-between">
            <CardTitle>Results</CardTitle>
            <div className="flex items-center gap-2">
              {videoGen.result.provider && <ProviderBadge provider={videoGen.result.provider} />}
              <Button variant="outline" size="sm" onClick={() => navigate('/media')}>
                View in Library
              </Button>
            </div>
          </CardHeader>
          <CardContent>
            <VideoGrid
              videos={videoGen.result.videos.map((vid) => ({
                src: vid.isUrl ? vid.data : `data:${vid.mimeType};base64,${vid.data}`,
                mimeType: vid.mimeType,
                alt: prompt,
              }))}
            />
          </CardContent>
        </Card>
      )}
    </div>
  );
}