ui/src/app/datasets/page.tsx

'use client';

import { useState } from 'react';
import { Modal } from '@/components/Modal';
import Link from 'next/link';
import { TextInput } from '@/components/formInputs';
import useDatasetList from '@/hooks/useDatasetList';
import { Button } from '@headlessui/react';
import { FaRegTrashAlt } from 'react-icons/fa';
import { openConfirm } from '@/components/ConfirmModal';
import { TopBar, MainContent } from '@/components/layout';
import UniversalTable, { TableColumn } from '@/components/UniversalTable';
import { apiClient } from '@/utils/api';
import { useRouter } from 'next/navigation';
import { usingBrowserDb } from '@/utils/env';
import { addUserDataset, removeUserDataset } from '@/utils/storage/datasetStorage';
import { useAuth } from '@/contexts/AuthContext';
import HFLoginButton from '@/components/HFLoginButton';

export default function Datasets() {
  const router = useRouter();
  const { datasets, status, refreshDatasets } = useDatasetList();
  const [newDatasetName, setNewDatasetName] = useState('');
  const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
  const { status: authStatus, namespace } = useAuth();
  const isAuthenticated = authStatus === 'authenticated';

  // Transform datasets array into rows with objects
  const tableRows = datasets.map(dataset => ({
    name: dataset,
    actions: dataset, // Pass full dataset name for actions
  }));

  const columns: TableColumn[] = [
    {
      title: 'Dataset Name',
      key: 'name',
      render: row => (
        <Link href={`/datasets/${row.name}`} className="text-gray-200 hover:text-gray-100">
          {row.name}
        </Link>
      ),
    },
    {
      title: 'Actions',
      key: 'actions',
      className: 'w-20 text-right',
      render: row => (
        <button
          className="text-gray-200 hover:bg-red-600 p-2 rounded-full transition-colors"
          onClick={() => handleDeleteDataset(row.name)}
        >
          <FaRegTrashAlt />
        </button>
      ),
    },
  ];

  const handleDeleteDataset = (datasetName: string) => {
    openConfirm({
      title: 'Delete Dataset',
      message: `Are you sure you want to delete the dataset "${datasetName}"? This action cannot be undone.`,
      type: 'warning',
      confirmText: 'Delete',
      onConfirm: () => {
        apiClient
          .post('/api/datasets/delete', { name: datasetName })
          .then(() => {
            console.log('Dataset deleted:', datasetName);
            if (usingBrowserDb) {
              removeUserDataset(datasetName);
            }
            refreshDatasets();
          })
          .catch(error => {
            console.error('Error deleting dataset:', error);
          });
      },
    });
  };

  const handleCreateDataset = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!isAuthenticated) {
      return;
    }
    try {
      const data = await apiClient
        .post('/api/datasets/create', { name: newDatasetName, namespace })
        .then(res => res.data);
      console.log('New dataset created:', data);
      if (usingBrowserDb && data?.name) {
        addUserDataset(data.name, data?.path || '');
      }
      refreshDatasets();
      setNewDatasetName('');
      setIsNewDatasetModalOpen(false);
    } catch (error) {
      console.error('Error creating new dataset:', error);
    }
  };

  const openNewDatasetModal = () => {
    if (!isAuthenticated) {
      return;
    }
    openConfirm({
      title: 'New Dataset',
      message: 'Enter the name of the new dataset:',
      type: 'info',
      confirmText: 'Create',
      inputTitle: 'Dataset Name',
      onConfirm: async (name?: string) => {
        if (!name) {
          console.error('Dataset name is required.');
          return;
        }
        if (!isAuthenticated) {
          return;
        }
        try {
          const data = await apiClient
            .post('/api/datasets/create', { name, namespace })
            .then(res => res.data);
          console.log('New dataset created:', data);
          if (usingBrowserDb && data?.name) {
            addUserDataset(data.name, data?.path || '');
          }
          if (data.name) {
            router.push(`/datasets/${data.name}`);
          } else {
            refreshDatasets();
          }
        } catch (error) {
          console.error('Error creating new dataset:', error);
        }
      },
    });
  };

  return (
    <>
      <TopBar>
        <div>
          <h1 className="text-2xl font-semibold text-gray-100">Datasets</h1>
        </div>
        <div className="flex-1"></div>
        <div>
          {isAuthenticated ? (
            <Button
              className="text-gray-200 bg-slate-600 px-4 py-2 rounded-md hover:bg-slate-500 transition-colors"
              onClick={() => openNewDatasetModal()}
            >
              New Dataset
            </Button>
          ) : (
            <span className="text-gray-600 bg-gray-900 px-3 py-1 rounded-md border border-gray-800">
              Sign in to add datasets
            </span>
          )}
        </div>
      </TopBar>

      <MainContent>
        {isAuthenticated ? (
          <UniversalTable
            columns={columns}
            rows={tableRows}
            isLoading={status === 'loading'}
            onRefresh={refreshDatasets}
          />
        ) : (
          <div className="border border-gray-800 rounded-lg p-6 bg-gray-900 text-gray-400 text-sm flex flex-col gap-4">
            <p>Sign in with Hugging Face or add an access token to manage datasets.</p>
            <div className="flex items-center gap-3">
              <HFLoginButton size="sm" />
              <Link href="/settings" className="text-xs text-blue-400 hover:text-blue-300">
                Manage authentication in Settings
              </Link>
            </div>
          </div>
        )}
      </MainContent>

      <Modal
        isOpen={isNewDatasetModalOpen}
        onClose={() => setIsNewDatasetModalOpen(false)}
        title="New Dataset"
        size="md"
      >
        <div className="space-y-4 text-gray-200">
          <form onSubmit={handleCreateDataset}>
            <div className="text-sm text-gray-400">
              This will create a new folder with the name below in your dataset folder.
            </div>
            <div className="mt-4">
              <TextInput label="Dataset Name" value={newDatasetName} onChange={value => setNewDatasetName(value)} />
            </div>

            <div className="mt-6 flex justify-end space-x-3">
              <button
                type="button"
                className="rounded-md bg-gray-700 px-4 py-2 text-gray-200 hover:bg-gray-600 focus:outline-none focus:ring-2 focus:ring-gray-500"
                onClick={() => setIsNewDatasetModalOpen(false)}
              >
                Cancel
              </button>
              <button
                type="submit"
                className="rounded-md bg-blue-600 px-4 py-2 text-white hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:opacity-50 disabled:cursor-not-allowed"
                disabled={!isAuthenticated}
              >
                Confirm
              </button>
            </div>
          </form>
        </div>
      </Modal>
    </>
  );
}