stable-diffusion.cpp-rest/webui/app/gallery/page.tsx

"use client";

import { useState, useEffect, useCallback } from "react";
import { Header } from "@/components/layout";
import { AppLayout } from "@/components/layout";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { apiClient, type JobDetailsResponse } from "@/lib/api";
import {
  Loader2,
  Download,
  RefreshCw,
  Calendar,
  Image as ImageIcon,
  X,
  ZoomIn,
} from "lucide-react";
import { downloadAuthenticatedImage } from "@/lib/utils";

interface GalleryImage {
  jobId: string;
  filename: string;
  url: string;
  thumbnailUrl: string;
  prompt?: string;
  negativePrompt?: string;
  width?: number;
  height?: number;
  steps?: number;
  cfgScale?: number;
  seed?: string;
  model?: string;
  createdAt: string;
  status: string;
}

interface ImageModalProps {
  image: GalleryImage | null;
  isOpen: boolean;
  onClose: () => void;
}

function ImageModal({ image, isOpen, onClose }: ImageModalProps) {
  if (!isOpen || !image) return null;

  const handleBackdropClick = (e: React.MouseEvent) => {
    if (e.target === e.currentTarget) {
      onClose();
    }
  };

  return (
    <div
      className="fixed inset-0 z-50 flex items-center justify-center bg-black/80 backdrop-blur-sm"
      onClick={handleBackdropClick}
    >
      <div className="relative max-w-[90vw] max-h-[90vh] flex flex-col">
        {/* Header with image info */}
        <div className="bg-background border-b border-border p-4 rounded-t-lg">
          <div className="flex items-center justify-between">
            <div className="flex items-center gap-2">
              <ImageIcon className="h-5 w-5" />
              <span className="font-medium">Generated Image</span>
              <Badge variant="secondary">{image.model || "Unknown"}</Badge>
            </div>
            <Button
              size="icon"
              variant="ghost"
              onClick={onClose}
              className="h-8 w-8"
            >
              <X className="h-4 w-4" />
            </Button>
          </div>

          {/* Image metadata */}
          <div className="mt-2 text-sm text-muted-foreground grid grid-cols-2 md:grid-cols-4 gap-2">
            {image.width && image.height && (
              <span>
                Size: {image.width}×{image.height}
              </span>
            )}
            {image.steps && <span>Steps: {image.steps}</span>}
            {image.cfgScale && <span>CFG: {image.cfgScale}</span>}
            {image.seed && <span>Seed: {image.seed}</span>}
          </div>

          {image.prompt && (
            <div className="mt-2">
              <p className="text-sm font-medium">Prompt:</p>
              <p className="text-sm text-muted-foreground line-clamp-2">
                {image.prompt}
              </p>
            </div>
          )}
        </div>

        {/* Image container */}
        <div className="flex items-center justify-center bg-muted/50 p-4 rounded-b-lg">
          <img
            src={image.url}
            alt="Generated image"
            className="max-w-full max-h-[70vh] object-contain rounded-lg shadow-lg"
          />
        </div>

        {/* Download button */}
        <div className="absolute bottom-4 right-4">
          <Button
            size="sm"
            onClick={() => {
              const authToken = localStorage.getItem("auth_token");
              const unixUser = localStorage.getItem("unix_user");
              downloadAuthenticatedImage(
                image.url,
                `gallery-${image.jobId}-${image.filename}`,
                authToken || undefined,
                unixUser || undefined,
              );
            }}
          >
            <Download className="h-4 w-4 mr-2" />
            Download
          </Button>
        </div>
      </div>
    </div>
  );
}

function GalleryGrid() {
  const [images, setImages] = useState<GalleryImage[]>([]);
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  const [selectedImage, setSelectedImage] = useState<GalleryImage | null>(null);
  const [isModalOpen, setIsModalOpen] = useState(false);
  const [sortBy, setSortBy] = useState<"newest" | "oldest">("newest");

  const loadGalleryImages = useCallback(async () => {
    try {
      setLoading(true);
      setError(null);

      console.log("Gallery: Loading images...");

      // Get queue status to find all completed jobs
      const queueStatus = await apiClient.getQueueStatus();
      console.log("Gallery: Queue status:", queueStatus);
      const completedJobs = queueStatus.jobs.filter(
        (job) => job.status === "completed",
      );
      console.log("Gallery: Completed jobs:", completedJobs);

      const galleryImages: GalleryImage[] = [];

      // Fetch individual job details to get output information
      for (const job of completedJobs) {
        const jobId = job.request_id || job.id || "";
        if (!jobId) continue;

        try {
          // Get detailed job information including outputs
          console.log(`Gallery: Fetching details for job ${jobId}`);
          const jobDetails: JobDetailsResponse =
            await apiClient.getJobStatus(jobId);
          console.log(`Gallery: Job ${jobId} details:`, jobDetails);

          // API response has outputs nested in job object
          if (
            jobDetails.job &&
            jobDetails.job.outputs &&
            jobDetails.job.outputs.length > 0
          ) {
            for (const output of jobDetails.job.outputs) {
              const filename = output.filename;
              const url = apiClient.getImageUrl(jobId, filename);

              // Create thumbnail URL (we'll use the same URL but let the browser handle scaling)
              const thumbnailUrl = url;

              galleryImages.push({
                jobId,
                filename,
                url,
                thumbnailUrl,
                prompt: jobDetails.job.prompt || "", // Get prompt from job details
                negativePrompt: "", // Job info doesn't include negative prompt
                width: undefined, // Job info doesn't include request details
                height: undefined, // Job info doesn't include request details
                steps: undefined, // Job info doesn't include request details
                cfgScale: undefined, // Job info doesn't include request details
                seed: undefined, // Job info doesn't include request details
                model: "Unknown", // Job info doesn't include request details
                createdAt:
                  jobDetails.job.created_at || new Date().toISOString(),
                status: jobDetails.job.status,
              });
            }
          }
        } catch (err) {
          console.warn(`Failed to fetch details for job ${jobId}:`, err);
          // Continue with other jobs even if one fails
        }
      }

      // Sort images
      galleryImages.sort((a, b) => {
        const dateA = new Date(a.createdAt).getTime();
        const dateB = new Date(b.createdAt).getTime();
        return sortBy === "newest" ? dateB - dateA : dateA - dateB;
      });

      console.log("Gallery: Final images array:", galleryImages);
      setImages(galleryImages);
    } catch (err) {
      console.error("Gallery: Error loading images:", err);
      setError(
        err instanceof Error ? err.message : "Failed to load gallery images",
      );
    } finally {
      setLoading(false);
    }
  }, [sortBy]);

  useEffect(() => {
    loadGalleryImages();
  }, [sortBy, loadGalleryImages]);

  const handleImageClick = (image: GalleryImage) => {
    setSelectedImage(image);
    setIsModalOpen(true);
  };

  const handleModalClose = () => {
    setIsModalOpen(false);
    setSelectedImage(null);
  };

  const handleDownload = (image: GalleryImage, e: React.MouseEvent) => {
    e.stopPropagation();
    const authToken = localStorage.getItem("auth_token");
    const unixUser = localStorage.getItem("unix_user");
    downloadAuthenticatedImage(
      image.url,
      `gallery-${image.jobId}-${image.filename}`,
      authToken || undefined,
      unixUser || undefined,
    );
  };

  const formatDate = (dateString: string) => {
    return new Date(dateString).toLocaleDateString("en-US", {
      year: "numeric",
      month: "short",
      day: "numeric",
      hour: "2-digit",
      minute: "2-digit",
    });
  };

  if (loading) {
    return (
      <AppLayout>
        <Header title="Gallery" description="Browse your generated images" />
        <div className="container mx-auto p-6">
          <div className="flex items-center justify-center h-96">
            <div className="flex items-center gap-2">
              <Loader2 className="h-6 w-6 animate-spin" />
              <span>Loading gallery...</span>
            </div>
          </div>
        </div>
      </AppLayout>
    );
  }

  if (error) {
    return (
      <AppLayout>
        <Header title="Gallery" description="Browse your generated images" />
        <div className="container mx-auto p-6">
          <div className="flex flex-col items-center justify-center h-96 gap-4">
            <div className="text-destructive text-center">
              <p className="text-lg font-medium">Error loading gallery</p>
              <p className="text-sm">{error}</p>
            </div>
            <Button onClick={loadGalleryImages} variant="outline">
              <RefreshCw className="h-4 w-4 mr-2" />
              Try Again
            </Button>
          </div>
        </div>
      </AppLayout>
    );
  }

  return (
    <AppLayout>
      <Header title="Gallery" description="Browse your generated images" />
      <div className="container mx-auto p-6">
        {/* Controls */}
        <div className="flex items-center justify-between mb-6">
          <div className="flex items-center gap-4">
            <h2 className="text-2xl font-bold">
              {images.length} {images.length === 1 ? "Image" : "Images"}
            </h2>
            <Button onClick={loadGalleryImages} variant="outline" size="sm">
              <RefreshCw className="h-4 w-4 mr-2" />
              Refresh
            </Button>
          </div>

          <div className="flex items-center gap-2">
            <span className="text-sm text-muted-foreground">Sort by:</span>
            <Button
              variant={sortBy === "newest" ? "default" : "outline"}
              size="sm"
              onClick={() => setSortBy("newest")}
            >
              Newest
            </Button>
            <Button
              variant={sortBy === "oldest" ? "default" : "outline"}
              size="sm"
              onClick={() => setSortBy("oldest")}
            >
              Oldest
            </Button>
          </div>
        </div>

        {/* Gallery Grid */}
        {images.length === 0 ? (
          <div className="flex flex-col items-center justify-center h-96 border-2 border-dashed border-border rounded-lg">
            <ImageIcon className="h-12 w-12 text-muted-foreground mb-4" />
            <h3 className="text-lg font-medium text-muted-foreground mb-2">
              No images found
            </h3>
            <p className="text-sm text-muted-foreground text-center max-w-md">
              Generate some images first using the Text to Image, Image to
              Image, or Inpainting tools.
            </p>
          </div>
        ) : (
          <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 lg:grid-cols-4 xl:grid-cols-5 gap-4">
            {images.map((image, index) => (
              <Card
                key={`${image.jobId}-${image.filename}-${index}`}
                className="group cursor-pointer overflow-hidden hover:shadow-lg transition-all duration-200"
                onClick={() => handleImageClick(image)}
              >
                <CardContent className="p-0">
                  <div className="relative aspect-square">
                    <img
                      src={image.thumbnailUrl}
                      alt={`Generated image ${index + 1}`}
                      className="w-full h-full object-cover"
                      loading="lazy"
                    />

                    {/* Overlay with actions */}
                    <div className="absolute inset-0 bg-black/50 opacity-0 group-hover:opacity-100 transition-opacity duration-200 flex items-center justify-center gap-2">
                      <Button
                        size="icon"
                        variant="secondary"
                        className="h-8 w-8"
                        onClick={(e) => {
                          e.stopPropagation();
                          handleImageClick(image);
                        }}
                        title="View full size"
                      >
                        <ZoomIn className="h-4 w-4" />
                      </Button>
                      <Button
                        size="icon"
                        variant="secondary"
                        className="h-8 w-8"
                        onClick={(e) => handleDownload(image, e)}
                        title="Download image"
                      >
                        <Download className="h-4 w-4" />
                      </Button>
                    </div>

                    {/* Date badge */}
                    <div className="absolute top-2 left-2">
                      <Badge variant="secondary" className="text-xs">
                        {formatDate(image.createdAt)}
                      </Badge>
                    </div>

                    {/* Model indicator */}
                    {image.model && image.model !== "Unknown" && (
                      <div className="absolute bottom-2 left-2">
                        <Badge
                          variant="outline"
                          className="text-xs max-w-20 truncate"
                        >
                          {image.model}
                        </Badge>
                      </div>
                    )}
                  </div>

                  {/* Image info */}
                  <div className="p-3">
                    <div className="flex items-center justify-between text-xs text-muted-foreground">
                      <span className="truncate">
                        {image.width && image.height
                          ? `${image.width}×${image.height}`
                          : "Unknown size"}
                      </span>
                      <div className="flex items-center gap-1">
                        <Calendar className="h-3 w-3" />
                        <span>
                          {new Date(image.createdAt).toLocaleDateString()}
                        </span>
                      </div>
                    </div>

                    {image.prompt && (
                      <p className="mt-1 text-xs text-muted-foreground line-clamp-2">
                        {image.prompt}
                      </p>
                    )}
                  </div>
                </CardContent>
              </Card>
            ))}
          </div>
        )}

        {/* Image Modal */}
        <ImageModal
          image={selectedImage}
          isOpen={isModalOpen}
          onClose={handleModalClose}
        />
      </div>
    </AppLayout>
  );
}

export default function GalleryPage() {
  return <GalleryGrid />;
}