ui/src/components/DatasetPreviewGrid.tsx

import { useCallback, useEffect, useMemo, useState } from 'react';
import { LuImageOff, LuLoader, LuBan } from 'react-icons/lu';
import DatasetImageCard from './DatasetImageCard';
import classNames from 'classnames';
import { apiClient } from '@/utils/api';

type LoadState = 'idle' | 'loading' | 'success' | 'error' | 'empty';

interface DatasetPreviewGridProps {
  datasetName?: string | null;
  refreshKey?: number;
  className?: string;
  emptyMessage?: string;
  compact?: boolean;
}

export default function DatasetPreviewGrid({
  datasetName,
  refreshKey = 0,
  className = '',
  emptyMessage = 'No files uploaded yet. Add images or videos to see them here.',
  compact = false,
}: DatasetPreviewGridProps) {
  const [status, setStatus] = useState<LoadState>('idle');
  const [items, setItems] = useState<{ img_path: string }[]>([]);

  const fetchMedia = useCallback(async () => {
    const trimmedName = datasetName?.trim();
    if (!trimmedName) {
      setItems([]);
      setStatus('idle');
      return;
    }

    setStatus('loading');
    try {
      const response = await apiClient
        .post('/api/datasets/listImages', { datasetName: trimmedName })
        .then(res => res.data);

      const sorted = Array.isArray(response?.images)
        ? [...response.images].sort((a, b) => a.img_path.localeCompare(b.img_path))
        : [];

      setItems(sorted);
      setStatus(sorted.length === 0 ? 'empty' : 'success');
    } catch (error) {
      console.error('Failed to fetch dataset preview:', error);
      setItems([]);
      setStatus('error');
    }
  }, [datasetName]);

  useEffect(() => {
    fetchMedia();
  }, [fetchMedia, refreshKey]);

  const content = useMemo(() => {
    switch (status) {
      case 'idle':
        return (
          <div className="text-sm text-gray-500 bg-gray-900/60 border border-gray-800 rounded-md px-4 py-6 text-center">
            Select or upload a dataset to preview its contents.
          </div>
        );
      case 'loading':
        return (
          <div className="flex flex-col items-center justify-center gap-3 py-8 text-sm text-gray-300">
            <LuLoader className="animate-spin w-6 h-6" />
            <span>Loading dataset contents…</span>
          </div>
        );
      case 'error':
        return (
          <div className="flex flex-col items-center justify-center gap-3 py-8 text-sm text-red-300 bg-red-950/30 border border-red-800 rounded-md">
            <LuBan className="w-6 h-6" />
            <span>Unable to load dataset media. Verify the dataset exists and try again.</span>
          </div>
        );
      case 'empty':
        return (
          <div className="flex flex-col items-center justify-center gap-3 py-8 text-sm text-gray-400 bg-gray-900/60 border border-dashed border-gray-700 rounded-md">
            <LuImageOff className="w-6 h-6" />
            <span>{emptyMessage}</span>
          </div>
        );
      case 'success':
        return (
          <div
            className={classNames(
              'grid gap-3 sm:gap-4',
              compact
                ? 'grid-cols-2 sm:grid-cols-3 lg:grid-cols-4 xl:grid-cols-5'
                : 'grid-cols-2 sm:grid-cols-3 md:grid-cols-4 xl:grid-cols-5 2xl:grid-cols-6',
            )}
          >
            {items.map(item => (
              <DatasetImageCard
                key={item.img_path}
                imageUrl={item.img_path}
                alt="Dataset media"
                onDelete={fetchMedia}
              />
            ))}
          </div>
        );
      default:
        return null;
    }
  }, [status, items, emptyMessage, compact, fetchMedia]);

  return <div className={classNames('space-y-4', className)}>{content}</div>;
}