import React, { useState, useMemo, useEffect, useCallback, lazy, Suspense } from 'react';
import { DatasetManager } from './components/DatasetManager';
import type { Dataset, DatasetMetadata } from './types';
import * as db from './services/dbService';
import * as apiService from './services/apiService';
import { Spinner } from './components/common/Spinner';

const ComparisonTool = lazy(() => import('./components/ComparisonTool').then(module => ({ default: module.ComparisonTool })));

type View = 'manager' | 'comparison';

const App: React.FC = () => {
  const [datasets, setDatasets] = useState<DatasetMetadata[]>([]);
  const [selectedDatasetId, setSelectedDatasetId] = useState<string | null>(null);
  const [activeDataset, setActiveDataset] = useState<Dataset | null>(null);
  const [view, setView] = useState<View>('manager');
  const [isLoading, setIsLoading] = useState(true);
  const [isNavigating, setIsNavigating] = useState(false);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    const loadInitialData = async () => {
      try {
        const localMeta = await db.getAllDatasetMetadata();
        let sharedMeta: DatasetMetadata[] = [];
        try {
            sharedMeta = await apiService.getSharedDatasetMetadata();
        } catch (e) {
            console.error("Could not load shared datasets, continuing with local.", e);
            setError("Could not load cloud datasets. The backend service may be unavailable. Local datasets are still accessible.");
        }
        
        const allMeta = [...sharedMeta, ...localMeta];

        setDatasets(allMeta);

        if (allMeta.length > 0) {
            // Select the most recent dataset by default
            const sortedMeta = [...allMeta].sort((a, b) => new Date(b.uploadDate).getTime() - new Date(a.uploadDate).getTime());
            setSelectedDatasetId(sortedMeta[0].id);
        }
      } catch (error) {
        console.error("Failed to load initial data", error);
        setError("A critical error occurred while loading local datasets.");
      } finally {
        setIsLoading(false);
      }
    };
    loadInitialData();
  }, []);

  const addDataset = async (newDataset: Dataset) => {
    await db.addDataset(newDataset);
    const localMeta = await db.getAllDatasetMetadata();
    const sharedMeta = datasets.filter(d => d.isShared); // Keep existing shared meta
    setDatasets([...sharedMeta, ...localMeta]);
    setSelectedDatasetId(newDataset.id);
  };

  const deleteDataset = async (id: string) => {
    await db.deleteDataset(id);
    setDatasets(prevDatasets => {
      const newDatasets = prevDatasets.filter(d => d.id !== id);
      if (selectedDatasetId === id) {
        const sortedMeta = [...newDatasets].sort((a, b) => new Date(b.uploadDate).getTime() - new Date(a.uploadDate).getTime());
        setSelectedDatasetId(sortedMeta.length > 0 ? sortedMeta[0].id : null);
      }
      return newDatasets;
    });
  };

  const renameDataset = async (id: string, newName: string) => {
    await db.renameDataset(id, newName);
    setDatasets(prev => prev.map(d => d.id === id ? { ...d, name: newName } : d));
  };
  
  const processedDatasets = useMemo(() => {
    return datasets.filter(d => d.processingState === 'processed');
  }, [datasets]);

  const getFullDataset = async (id: string): Promise<Dataset | null> => {
    const meta = datasets.find(d => d.id === id);
    if (!meta) return null;

    if (meta.isShared) {
      return apiService.getSharedDataset(id);
    } else {
      return db.getDataset(id);
    }
  };

  const handleOpenComparisonTool = useCallback(async () => {
    if (!selectedDatasetId) return;
    const selectedMeta = datasets.find(d => d.id === selectedDatasetId);
    if (!selectedMeta || selectedMeta.processingState !== 'processed') return;

    setView('comparison');
    setActiveDataset(null);
    setIsNavigating(true);

    try {
        const fullDataset = await getFullDataset(selectedDatasetId);
        if (!fullDataset) {
            throw new Error(`Failed to load dataset ${selectedDatasetId}.`);
        }

        // *** NEW LOGIC ***
        // If it's a local dataset, ensure it's in the backend's cache before proceeding.
        if (!fullDataset.isShared) {
            console.log("Local dataset selected. Ensuring it's cached on the backend...");
            await apiService.ensureDatasetInCache(fullDataset);
            console.log("Backend cache confirmed.");
        }

        setActiveDataset(fullDataset);

    } catch (error) {
        console.error("Error preparing comparison tool:", error);
        alert(`Error: Could not load the selected dataset. ${error instanceof Error ? error.message : ''}`);
        setView('manager'); // Go back on error
    } finally {
        setIsNavigating(false);
    }
  }, [selectedDatasetId, datasets]);

  const handleDatasetChange = useCallback(async (newId: string) => {
      setSelectedDatasetId(newId);
      
      setActiveDataset(null);
      setIsNavigating(true);

      try {
        const fullDataset = await getFullDataset(newId);
        if (!fullDataset) {
            throw new Error(`Failed to load dataset ${newId}.`);
        }

        // Also ensure cache is hydrated when switching datasets inside the tool
        if (!fullDataset.isShared) {
            await apiService.ensureDatasetInCache(fullDataset);
        }
        
        setActiveDataset(fullDataset);

      } catch (error) {
            console.error(`Error switching dataset to ${newId}:`, error);
            setActiveDataset(null);
      } finally {
          setIsNavigating(false);
      }
  }, []);


  const mainContent = () => {
    if (isLoading) {
        return <div className="flex justify-center items-center h-64"><Spinner /><span>Loading Datasets...</span></div>;
    }
    
    const errorBanner = error ? (
        <div className="bg-red-800/50 border border-red-600 text-red-200 px-4 py-3 rounded-lg mb-4" role="alert">
          <p>
            <strong className="font-bold">Cloud Connection Error:</strong> {error}
          </p>
        </div>
    ) : null;

    if (view === 'manager') {
      return (
        <DatasetManager
          datasets={datasets}
          selectedDatasetId={selectedDatasetId}
          onSelectDataset={setSelectedDatasetId}
          onAddDataset={addDataset}
          onDeleteDataset={deleteDataset}
          onRenameDataset={renameDataset}
          onOpenComparisonTool={handleOpenComparisonTool}
          onGetFullDataset={getFullDataset}
          errorBanner={errorBanner}
        />
      );
    }
    
    if (view === 'comparison') {
        const fallbackUI = <div className="flex justify-center items-center h-64"><Spinner /><span>Loading Comparison Tool...</span></div>;
        if (isNavigating || !activeDataset) {
            return fallbackUI;
        }
        return (
            <Suspense fallback={fallbackUI}>
                <ComparisonTool
                    key={activeDataset.id}
                    dataset={activeDataset}
                    allDatasets={processedDatasets}
                    onDatasetChange={handleDatasetChange}
                    onBack={() => { setView('manager'); setActiveDataset(null); }}
                />
            </Suspense>
        );
    }
    return null;
  };

  return (
    <div className="min-h-screen bg-gray-900 text-gray-200 font-sans">
      <header className="bg-gray-800/50 backdrop-blur-sm border-b border-gray-700 sticky top-0 z-20">
        <div className="container mx-auto px-4 sm:px-6 lg:px-8 py-4 flex items-center justify-between">
          <h1 className="text-2xl font-bold text-cyan-400">
            Cross-Modal Object Comparison Tool
          </h1>
        </div>
      </header>
      <main className="container mx-auto px-4 sm:px-6 lg:px-8 py-8">
        {mainContent()}
      </main>
    </div>
  );
};

export default App;