ui-nyla/src/components/Dashboard/DashboardChat/useWorkflowManager.ts

import { useState, useEffect, useCallback, useRef, useMemo } from 'react';
import { useWorkflow, useWorkflowStatus, useWorkflowMessages, useWorkflowLogs, useWorkflowOperations, StartWorkflowRequest } from '../../../hooks/useWorkflows';
import { WorkflowState, WorkflowActions } from './dashboardChatAreaTypes';

export function useWorkflowManager(initialWorkflowId?: string | null): [WorkflowState, WorkflowActions] {
  // Core state
  const [currentWorkflowId, setCurrentWorkflowId] = useState<string | null>(initialWorkflowId || null);
  const [isPolling, setIsPolling] = useState(false);
  const [pendingMessages, setPendingMessages] = useState<any[]>([]);
  const [sentUserMessages, setSentUserMessages] = useState<Set<string>>(new Set()); // Track sent user messages
  const [selectedPrompt, setSelectedPrompt] = useState<any | null>(null); // Selected prompt state
  const pollingIntervalRef = useRef<number | null>(null);

  // Hook-based data fetching
  const { workflow, loading: workflowLoading, error: workflowError } = useWorkflow(currentWorkflowId);
  const { status: workflowStatus, loading: statusLoading, error: statusError, refetch: refetchStatus } = useWorkflowStatus(currentWorkflowId);
  const { messages, loading: messagesLoading, error: messagesError, refetch: refetchMessages } = useWorkflowMessages(currentWorkflowId);
  const { logs, loading: logsLoading, error: logsError, refetch: refetchLogs } = useWorkflowLogs(currentWorkflowId);
  const { startWorkflow, stopWorkflow: stopWorkflowRequest } = useWorkflowOperations();

  // Use status for real-time updates, fallback to workflow for initial data
  const currentWorkflow = workflowStatus || workflow;

  // Helper to create optimistic user message
  const createOptimisticMessage = useCallback((prompt: string, fileIds: string[] = []) => {
    // Use UTC timestamp in seconds (float) to match backend expectation
    const timestamp = Math.floor(Date.now() / 1000);
    return {
      id: `temp-${Date.now()}`,
      workflowId: currentWorkflowId || 'pending',
      message: prompt,
      role: 'user' as const,
      status: 'pending',
      sequenceNr: 0,
      publishedAt: timestamp,
      success: true,
      fileIds: fileIds.length > 0 ? fileIds : undefined
    };
  }, [currentWorkflowId]);

  // Combined loading and error states
  const isLoading = workflowLoading || statusLoading || messagesLoading || logsLoading;
  const error = workflowError || statusError || messagesError || logsError;

  // Filter out ALL user messages from backend - we only show user messages from pendingMessages
  const filteredMessages = useMemo(() => {
    if (!messages) return [];
    
    // If we've tracked ANY sent messages, always filter out user messages from backend
    // This prevents flickering during new workflow creation
    if (sentUserMessages.size > 0) {
      return messages.filter(msg => msg.role !== 'user');
    }
    
    // If no tracked sent messages, this is a historical workflow - show all messages
    return messages;
  }, [messages, sentUserMessages]);

  // Auto-polling for active workflows and message updates
  useEffect(() => {
    if (isPolling && currentWorkflowId) {
      pollingIntervalRef.current = window.setInterval(() => {
        refetchStatus();
        if (currentWorkflow) {
          const isActive = ['running', 'processing', 'started'].includes(currentWorkflow.status);
          if (isActive) {
            refetchMessages();
            refetchLogs();
          }
        }
      }, 2000);
    }

    return () => {
      if (pollingIntervalRef.current) {
        clearInterval(pollingIntervalRef.current);
        pollingIntervalRef.current = null;
      }
    };
  }, [isPolling, currentWorkflowId, currentWorkflow?.status, refetchStatus, refetchMessages, refetchLogs]);

  // Actions
  const loadWorkflow = useCallback(async (workflowId: string) => {
    setCurrentWorkflowId(workflowId);
  }, []);

  const startNewWorkflow = useCallback(async (prompt: string, fileIds: string[] = []): Promise<string | null> => {
    // Add optimistic message immediately
    const optimisticMessage = createOptimisticMessage(prompt, fileIds);
    setPendingMessages(prev => [...prev, optimisticMessage]);
    
    // Track this message as sent
    setSentUserMessages(prev => new Set(prev).add(prompt.trim()));

    const workflowData: StartWorkflowRequest = {
      prompt,
      listFileId: fileIds
    };

    const result = await startWorkflow(workflowData);
    
    if (result.success && result.data) {
      const newWorkflowId = result.data.id;
      setCurrentWorkflowId(newWorkflowId);
      setIsPolling(true);
      setTimeout(() => {
        refetchMessages();
        refetchLogs();
      }, 500);
      return newWorkflowId;
    } else {
      // Remove optimistic message and sent tracking on failure
      setPendingMessages(prev => prev.filter(msg => msg.id !== optimisticMessage.id));
      setSentUserMessages(prev => {
        const newSet = new Set(prev);
        newSet.delete(prompt.trim());
        return newSet;
      });
    }
    return null;
  }, [startWorkflow, refetchMessages, createOptimisticMessage]);

  const continueWorkflow = useCallback(async (prompt: string, fileIds: string[] = []): Promise<boolean> => {
    if (!currentWorkflowId) return false;

    // Add optimistic message immediately
    const optimisticMessage = createOptimisticMessage(prompt, fileIds);
    setPendingMessages(prev => [...prev, optimisticMessage]);
    
    // Track this message as sent
    setSentUserMessages(prev => new Set(prev).add(prompt.trim()));

    const workflowData: StartWorkflowRequest = {
      prompt,
      listFileId: fileIds
    };

    const result = await startWorkflow(workflowData, currentWorkflowId);
    
    if (result.success) {
      setIsPolling(true);
      setTimeout(() => {
        refetchMessages();
        refetchLogs();
      }, 500);
      return true;
    } else {
      // Remove optimistic message and sent tracking on failure
      setPendingMessages(prev => prev.filter(msg => msg.id !== optimisticMessage.id));
      setSentUserMessages(prev => {
        const newSet = new Set(prev);
        newSet.delete(prompt.trim());
        return newSet;
      });
    }
    return false;
  }, [currentWorkflowId, startWorkflow, refetchMessages, refetchLogs, createOptimisticMessage]);

  const stopWorkflow = useCallback(async (): Promise<boolean> => {
    if (!currentWorkflowId) return false;

    const result = await stopWorkflowRequest(currentWorkflowId);
    if (result) {
      // Immediately refresh status to reflect the stopped state
      setTimeout(() => {
        refetchStatus();
        refetchMessages();
        refetchLogs();
      }, 500);
      return true;
    }
    return false;
  }, [currentWorkflowId, stopWorkflowRequest, refetchStatus, refetchMessages, refetchLogs]);

  const clearWorkflow = useCallback(() => {
    setCurrentWorkflowId(null);
    setIsPolling(false);
    setPendingMessages([]);
    setSentUserMessages(new Set());
  }, []);

  const selectPrompt = useCallback((prompt: any | null) => {
    setSelectedPrompt(prompt);
  }, []);

  const clearPrompt = useCallback(() => {
    setSelectedPrompt(null);
  }, []);

  // Only clear pending messages when workflow changes to a different ID
  // (not when creating a new workflow)
  const previousWorkflowId = useRef(currentWorkflowId);
  useEffect(() => {
    const prev = previousWorkflowId.current;
    const current = currentWorkflowId;
    
    // If we're switching between different existing workflows, clear state
    if (prev && current && prev !== current) {
      setPendingMessages([]);
      setSentUserMessages(new Set());
    }
    // If we're going from a workflow to no workflow, clear state
    else if (prev && !current) {
      setPendingMessages([]);
      setSentUserMessages(new Set());
    }
    // If we're going from no workflow to a workflow (new workflow creation), keep pending messages
    
    previousWorkflowId.current = current;
  }, [currentWorkflowId]);

  // Sync with external workflow ID changes
  useEffect(() => {
    if (initialWorkflowId !== currentWorkflowId) {
      if (initialWorkflowId) {
        loadWorkflow(initialWorkflowId);
      } else {
        setCurrentWorkflowId(null);
        setIsPolling(false);
      }
    }
  }, [initialWorkflowId]);

  // Auto-enable polling only for active workflows
  useEffect(() => {
    if (currentWorkflowId && currentWorkflow) {
      const isActive = ['running', 'processing', 'started'].includes(currentWorkflow.status);
      setIsPolling(isActive);
    } else {
      setIsPolling(false);
    }
  }, [currentWorkflowId, currentWorkflow?.status]);

  const state: WorkflowState = {
    currentWorkflowId,
    workflow: currentWorkflow,
    messages: filteredMessages,
    pendingMessages,
    logs: logs || [],
    isLoading,
    error,
    selectedPrompt
  };

  const actions: WorkflowActions = useMemo(() => ({
    loadWorkflow,
    startNewWorkflow,
    continueWorkflow,
    stopWorkflow,
    clearWorkflow,
    selectPrompt,
    clearPrompt
  }), [loadWorkflow, startNewWorkflow, continueWorkflow, stopWorkflow, clearWorkflow, selectPrompt, clearPrompt]);

  return [state, actions];
}