semblance/src/components/ui/BulkExportProgressModal.tsx

import React, { useState, useEffect, useRef } from 'react';
import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from '@/components/ui/dialog';
import { Progress } from '@/components/ui/progress';
import { Button } from '@/components/ui/button';
import { X, Download, CheckCircle, AlertCircle } from 'lucide-react';
import { cn } from '@/lib/utils';
import { toast } from 'sonner';

interface BulkExportProgressData {
  task_id: string;
  task_type: 'bulk_persona_export';
  progress: number;
  current_item: string;
  completed_count: number;
  total_count: number;
  current_persona_name?: string;
}

interface BulkExportProgressModalProps {
  isOpen: boolean;
  onClose: () => void;
  taskId: string | null;
  exportFormat: 'markdown' | 'json' | 'csv';
  personaCount: number;
  onCancel?: () => Promise<boolean>;
  onDownload?: (filePath: string) => void;
}

export const BulkExportProgressModal: React.FC<BulkExportProgressModalProps> = ({
  isOpen,
  onClose,
  taskId,
  exportFormat,
  personaCount,
  onCancel,
  onDownload
}) => {
  const [progress, setProgress] = useState(0);
  const [currentItem, setCurrentItem] = useState('Initializing export...');
  const [completedCount, setCompletedCount] = useState(0);
  const [currentPersonaName, setCurrentPersonaName] = useState<string | null>(null);
  const [isComplete, setIsComplete] = useState(false);
  const [hasError, setHasError] = useState(false);
  const [isCancelling, setIsCancelling] = useState(false);
  const [errorMessage, setErrorMessage] = useState<string | null>(null);
  const [downloadFilePath, setDownloadFilePath] = useState<string | null>(null);
  
  const stateRef = useRef({ taskId, isComplete, hasError });
  stateRef.current = { taskId, isComplete, hasError };

  // Reset state when modal opens
  useEffect(() => {
    if (isOpen && taskId) {
      setProgress(0);
      setCurrentItem('Initializing export...');
      setCompletedCount(0);
      setCurrentPersonaName(null);
      setIsComplete(false);
      setHasError(false);
      setIsCancelling(false);
      setErrorMessage(null);
      setDownloadFilePath(null);
    }
  }, [isOpen, taskId]);

  // Set up WebSocket event listeners for bulk export progress
  useEffect(() => {
    const handleBulkExportProgress = (event: CustomEvent) => {
      const data: BulkExportProgressData = event.detail;
      
      // Only handle events for our task
      if (data.task_id !== stateRef.current.taskId) {
        return;
      }

      setProgress(data.progress);
      setCurrentItem(data.current_item);
      setCompletedCount(data.completed_count);
      setCurrentPersonaName(data.current_persona_name || null);
    };

    const handleTaskCompleted = (event: CustomEvent) => {
      const data = event.detail;
      
      if (data.task_id !== stateRef.current.taskId) {
        return;
      }

      setIsComplete(true);
      setProgress(100);
      setCurrentItem('Export completed successfully!');
      
      // Store the file path for download
      if (data.file_path) {
        setDownloadFilePath(data.file_path);
      }
      
      toast.success(`Successfully exported ${personaCount} persona profiles!`);
    };

    const handleTaskFailed = (event: CustomEvent) => {
      const data = event.detail;
      
      if (data.task_id !== stateRef.current.taskId) {
        return;
      }

      setHasError(true);
      setErrorMessage(data.message || 'Export failed');
      setCurrentItem('Export failed');
      toast.error('Export failed: ' + (data.message || 'Unknown error'));
    };

    const handleTaskCancelled = (event: CustomEvent) => {
      const data = event.detail;
      
      if (data.task_id !== stateRef.current.taskId) {
        return;
      }

      setIsCancelling(false);
      setCurrentItem('Export cancelled');
      toast.success('Export cancelled successfully');
      
      // Close modal after short delay
      setTimeout(() => {
        onClose();
      }, 1000);
    };

    // Register window event listeners
    window.addEventListener('ws:bulk_export_progress', handleBulkExportProgress as EventListener);
    window.addEventListener('ws:task_completed', handleTaskCompleted as EventListener);
    window.addEventListener('ws:task_failed', handleTaskFailed as EventListener);
    window.addEventListener('ws:task_cancelled', handleTaskCancelled as EventListener);

    // Cleanup listeners
    return () => {
      window.removeEventListener('ws:bulk_export_progress', handleBulkExportProgress as EventListener);
      window.removeEventListener('ws:task_completed', handleTaskCompleted as EventListener);
      window.removeEventListener('ws:task_failed', handleTaskFailed as EventListener);
      window.removeEventListener('ws:task_cancelled', handleTaskCancelled as EventListener);
    };
  }, []); // Only set up once

  const handleCancel = async () => {
    if (!taskId || isCancelling) return;
    
    setIsCancelling(true);
    setCurrentItem('Cancelling export...');
    
    const success = onCancel ? await onCancel() : false;
    
    if (!success) {
      setIsCancelling(false);
    }
  };

  const handleDownload = () => {
    if (downloadFilePath && onDownload) {
      onDownload(downloadFilePath);
    }
  };

  const handleClose = () => {
    if (isComplete || hasError) {
      onClose();
    }
  };

  const getFormatDisplay = () => {
    switch (exportFormat) {
      case 'markdown': return 'Markdown';
      case 'json': return 'JSON';
      case 'csv': return 'CSV';
      default: return 'Files';
    }
  };

  const getProgressColor = () => {
    if (hasError) return 'bg-red-500';
    if (isComplete) return 'bg-green-500';
    if (isCancelling) return 'bg-orange-500';
    return 'bg-blue-500';
  };

  return (
    <Dialog open={isOpen} onOpenChange={handleClose}>
      <DialogContent className="sm:max-w-md">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            {hasError ? (
              <AlertCircle className="h-5 w-5 text-red-500" />
            ) : isComplete ? (
              <CheckCircle className="h-5 w-5 text-green-500" />
            ) : (
              <Download className="h-5 w-5 text-blue-500" />
            )}
            Bulk Export - {getFormatDisplay()}
          </DialogTitle>
          <DialogDescription>
            Exporting {personaCount} persona{personaCount !== 1 ? 's' : ''} to {getFormatDisplay().toLowerCase()} format
          </DialogDescription>
        </DialogHeader>

        <div className="space-y-4">
          {/* Progress Bar */}
          <div className="space-y-2">
            <Progress 
              value={progress} 
              className={cn(
                "w-full transition-all duration-200",
                hasError && "opacity-75",
                isCancelling && "opacity-60"
              )}
            />
            <div className="flex justify-between text-sm text-muted-foreground">
              <span>{Math.round(progress)}%</span>
              <span>{completedCount}/{personaCount}</span>
            </div>
          </div>

          {/* Current Status */}
          <div className="space-y-1">
            <div className="text-sm font-medium">
              {currentItem}
            </div>
            {currentPersonaName && !isComplete && !hasError && (
              <div className="text-sm text-muted-foreground">
                Processing: {currentPersonaName}
              </div>
            )}
          </div>

          {/* Error Message */}
          {hasError && errorMessage && (
            <div className="p-3 bg-red-50 border border-red-200 rounded-md">
              <div className="text-sm text-red-800">
                {errorMessage}
              </div>
            </div>
          )}

          {/* Action Buttons */}
          <div className="flex justify-between">
            <div>
              {/* Cancel Button */}
              {!isComplete && !hasError && taskId && (
                <Button
                  variant="outline"
                  onClick={handleCancel}
                  disabled={isCancelling}
                  className="text-muted-foreground hover:text-destructive hover:border-destructive"
                >
                  {isCancelling ? 'Cancelling...' : 'Cancel'}
                </Button>
              )}
            </div>
            
            <div className="flex gap-2">
              {/* Download Button */}
              {isComplete && downloadFilePath && (
                <Button
                  onClick={handleDownload}
                  className="bg-green-600 hover:bg-green-700"
                >
                  <Download className="h-4 w-4 mr-2" />
                  Download
                </Button>
              )}
              
              {/* Close Button */}
              {(isComplete || hasError) && (
                <Button
                  variant="outline"
                  onClick={handleClose}
                >
                  Close
                </Button>
              )}
            </div>
          </div>

          {/* Success Message */}
          {isComplete && (
            <div className="p-3 bg-green-50 border border-green-200 rounded-md">
              <div className="text-sm text-green-800">
                Export completed successfully! Your {getFormatDisplay().toLowerCase()} files are ready for download.
              </div>
            </div>
          )}
        </div>
      </DialogContent>
    </Dialog>
  );
};

export default BulkExportProgressModal;