import {createContext, useCallback, useContext, useEffect, useMemo, useRef, useState} from "@wordpress/element";
import {__, sprintf} from "@wordpress/i18n";
import {handleError, handleSuccess, handleWarning} from "../../../helpers/notifications";
import {
    DATASET_PROCESSES_INFO as SOURCES_DATASET_PROCESSES_INFO
} from "../../pages/settings/pages/knowledge/sources/_data";
import {
    DATASET_PROCESSES_INFO as RECOMMENDATIONS_DATASET_PROCESSES_INFO
} from "../../pages/settings/pages/knowledge/recommendations/_data";
import {delay} from "../../../../components/chatbots/includes/helpers";
import {parseDateString} from "../../../../helpers";
import useProgressDocumentTitle from "../../../../helpers/hooks/use-progress-document-title";
import {CancelJob, CreateJob, GetJob, PauseJob, ProcessJob, ResumeJob} from "../../../rest/jobs";

// Job status constants
export const JOB_STATUS_PENDING = 'pending';
export const JOB_STATUS_GENERATING_TASKS = 'generating_tasks';
export const JOB_STATUS_PROCESSING = 'processing';
export const JOB_STATUS_PAUSED = 'paused';
export const JOB_STATUS_COMPLETED = 'completed';
export const JOB_STATUS_FAILED = 'failed';
export const JOB_STATUS_CANCELED = 'canceled';

// Polling interval constants (in milliseconds)
const JOB_POLLING_INTERVAL = 500;

// Process status mapping
const JOB_TO_PROCESS_STATUS = {
    [JOB_STATUS_PENDING]: 'start',
    [JOB_STATUS_GENERATING_TASKS]: 'start',
    [JOB_STATUS_PROCESSING]: 'start',
    [JOB_STATUS_PAUSED]: 'pause',
    [JOB_STATUS_COMPLETED]: 'complete',
    [JOB_STATUS_FAILED]: 'cancel',
    [JOB_STATUS_CANCELED]: 'cancel',
};

// Create the context
const BackgroundProcessContext = createContext({
    // Process state (now job-based)
    processState: null,
    activeJobs: [], // Array of active job objects

    // Progress states
    generatingProgress: false,
    syncingProgress: false,
    deletingProgress: false,

    // Loading states
    saving: false,
    pausing: false,
    canceling: false,
    preparing: false,
    syncing: [],
    deleting: [],

    // Process management functions
    checkActiveProcess: () => true,
    pauseProcess: () => {
    },
    cancelProcess: () => {
    },
    createJob: () => {
    },
    followJob: () => {
    },
    processJobSequentially: () => {
    },

    // Progress update functions
    setGeneratingProgress: () => {
    },
    setSyncingProgress: () => {
    },
    setDeletingProgress: () => {
    },

    // Loading state functions
    setSaving: () => {
    },
    setPausing: () => {
    },
    setCanceling: () => {
    },
    setPreparing: () => {
    },
    setSyncing: () => {
    },
    setDeleting: () => {
    },

    // Internal state setters
    setProcessState: () => {
    },
    setActiveJobs: () => {
    },
});

// Provider component
export function BackgroundProcessProvider({children, notifications, pageSlug}) {
    // Process state - tracks the primary job being displayed
    const [processState, setProcessState] = useState(null);

    // Active jobs - all jobs running on this page
    const [activeJobs, setActiveJobs] = useState([]);

    // Progress states
    const [generatingProgress, setGeneratingProgress] = useState(false);
    const [syncingProgress, setSyncingProgress] = useState(false);
    const [deletingProgress, setDeletingProgress] = useState(false);

    // Loading states
    const [saving, setSaving] = useState(false);
    const [pausing, setPausing] = useState(false);
    const [canceling, setCanceling] = useState(false);
    const [preparing, setPreparing] = useState(false);
    const [syncing, setSyncing] = useState([]);
    const [deleting, setDeleting] = useState([]);

    // Track processed failed job reasons to avoid duplicates during process loop
    const processedFailedReasons = useRef(new Set());

    // Track if the current process was manually paused/canceled to avoid duplicate notifications
    const isManualActionRef = useRef(false);

    // Stop flags for polling and processing
    const stopPollingRef = useRef(new Map()); // Map of jobId -> stop flag
    const stopProcessingRef = useRef(new Map()); // Map of jobId -> stop flag
    const pollingInProgressRef = useRef(new Map()); // Map of jobId -> boolean (request in progress)
    const processingInProgressRef = useRef(new Map()); // Map of jobId -> boolean
    const resumeInProgressRef = useRef(new Map()); // Map of jobId -> boolean to prevent duplicate resume calls
    const pollingLoopActiveRef = useRef(new Map()); // Map of jobId -> boolean (polling loop is running)

    // Map to track onComplete callbacks for jobs: jobId -> onComplete callback
    const onCompleteCallbacksRef = useRef(new Map());

    // Map to track resetProgress callbacks for jobs: jobId -> resetProgress callback
    const resetProgressCallbacksRef = useRef(new Map());

    // Map to track onProgressUpdate callbacks for jobs: jobId -> onProgressUpdate callback
    const onProgressUpdateCallbacksRef = useRef(new Map());

    // Map to track if completion notification was shown for jobs: jobId -> boolean
    const completionNotificationShownRef = useRef(new Map());

    // Map to track job types for jobs: jobId -> jobType
    const jobTypesRef = useRef(new Map());

    // Map to track statusMessages for jobs: jobId -> statusMessages object
    const statusMessagesRef = useRef(new Map());

    // Map to track previous progress for jobs: jobId -> previousProgress (to avoid unnecessary dataset updates)
    const previousProgressRef = useRef(new Map());

    // Map to track previous completed_parents_count for jobs: jobId -> previousCompletedParentsCount
    // Used to determine when to update datasets list (only when count changes)
    const previousCompletedParentsCountRef = useRef(new Map());

    // Map to track previous errors length for jobs: jobId -> previousErrorsLength
    // Used to determine when to update datasets list (only when errors length changes)
    const previousErrorsLengthRef = useRef(new Map());

    // Ref to track latest processState for reliable access in async functions
    const processStateRef = useRef(null);

    // Keep processStateRef in sync with processState
    useEffect(() => {
        processStateRef.current = processState;
    }, [processState]);

    const {updateTitle: updateDocumentProgressTitle, restoreTitle: restoreDocumentTitle} = useProgressDocumentTitle();

    /**
     * Get the correct DATASET_PROCESSES_INFO based on pageSlug
     * Onboarding uses sources, recommendations uses recommendations, sources uses sources
     *
     * @return {object} DATASET_PROCESSES_INFO object
     */
    const getDatasetProcessesInfo = useCallback(() => {
        // Onboarding uses sources
        if (pageSlug === 'recommendations') {
            return RECOMMENDATIONS_DATASET_PROCESSES_INFO;
        }
        // Default to sources (for 'sources' and 'onboarding')
        return SOURCES_DATASET_PROCESSES_INFO;
    }, [pageSlug]);

    /**
     * Helper function to update document title with progress
     *
     * @param {number} progress Progress percentage (0-100)
     * @param {string} processType Process type (e.g., 'dataset_generating', 'dataset_sync')
     */
    const updateDocumentTitleWithProgress = useCallback((progress, processType) => {
        if (pageSlug !== 'onboarding') {
            updateDocumentProgressTitle(progress, (clampedProgress, originalTitle) => {
                const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                const processTitle = DATASET_PROCESSES_INFO[processType]?.inPageTitle || DATASET_PROCESSES_INFO[processType]?.processName || processType;
                return `${clampedProgress}% ${processTitle} - ${originalTitle}`;
            }, {resetOnInvalid: false});
        }
    }, [updateDocumentProgressTitle, pageSlug, getDatasetProcessesInfo]);

    /**
     * Helper function to get the appropriate progress setter based on job type
     *
     * @param {string} jobType Job type (e.g., 'dataset_generating', 'dataset_sync', 'dataset_delete')
     * @return {function} Progress setter function
     */
    const getProgressSetter = useCallback((jobType) => {
        return jobType === 'dataset_generating' ? setGeneratingProgress :
            jobType === 'dataset_sync' ? setSyncingProgress :
                jobType === 'dataset_delete' ? setDeletingProgress :
                    () => {
                    };
    }, [setGeneratingProgress, setSyncingProgress, setDeletingProgress]);

    useEffect(() => {
        // Initialize stop flags
        return () => {
            // Cleanup: stop all polling and processing
            stopPollingRef.current.forEach((_, jobId) => {
                stopPollingRef.current.set(jobId, true);
            });
            stopProcessingRef.current.forEach((_, jobId) => {
                stopProcessingRef.current.set(jobId, true);
            });
            restoreDocumentTitle();
        }
    }, []);

    /**
     * Check if another process is active on this page
     * Only one process per page is allowed
     */
    const checkActiveProcess = useCallback((currentProcessType) => {
        // Check if there's an active job on this page
        const hasActiveJob = activeJobs.some(job => {
            // Skip jobs that don't have a status (e.g., completed jobs that were removed)
            if (!job.status || job.completed === true) {
                return false;
            }
            const jobStatus = job.status;
            const isActive = jobStatus === JOB_STATUS_PROCESSING ||
                jobStatus === JOB_STATUS_GENERATING_TASKS ||
                jobStatus === JOB_STATUS_PENDING;
            return isActive && job.type !== currentProcessType;
        });

        if (hasActiveJob) {
            const activeJob = activeJobs.find(job => {
                // Skip jobs that don't have a status (e.g., completed jobs that were removed)
                if (!job.status || job.completed === true) {
                    return false;
                }
                const jobStatus = job.status;
                return (jobStatus === JOB_STATUS_PROCESSING ||
                        jobStatus === JOB_STATUS_GENERATING_TASKS ||
                        jobStatus === JOB_STATUS_PENDING) &&
                    job.type !== currentProcessType;
            });
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            handleWarning({
                message: "There is an active process on this page.",
                data: activeJob,
            }, notifications.set, {
                title: sprintf(__("%s process is already running.", 'limb-chatbot'), DATASET_PROCESSES_INFO[activeJob?.type]?.processName || activeJob?.type),
                description: __("Please wait for the current process to complete.", 'limb-chatbot'),
            });
            return false;
        }
        return true;
    }, [activeJobs, notifications]);

    /**
     * Handle job errors by showing error notifications
     *
     * @param {Array} errors Array of error objects
     * @param {string} processType Type of the process
     */
    const handleJobErrors = useCallback((errors, processType) => {
        if (!errors || !Array.isArray(errors) || errors.length === 0) {
            return;
        }

        // Filter out errors that have already been processed
        const newErrors = errors.filter(error => {
            if (!error) {
                return false;
            }

            // Use error message as identifier
            const errorKey = typeof error === 'string' ? error : JSON.stringify(error);

            // Check if this error has already been processed
            if (processedFailedReasons.current.has(errorKey)) {
                return false;
            }

            // Add to processed errors
            processedFailedReasons.current.add(errorKey);
            return true;
        });

        // Show error notification for each unique error
        newErrors.forEach((error) => {
            const errorMessage = typeof error === 'string' ? error : (error.message || JSON.stringify(error));

            // Check if this is a warning (based on error structure)
            const isWarning = typeof error === 'object' && (
                error.warning === true ||
                error.type === 'warning' ||
                error.error_data?.warning === true
            );

            const handler = isWarning ? handleWarning : handleError;
            handler(false, notifications.set, {
                title: errorMessage,
            });
        });
    }, [notifications]);

    /**
     * Poll job status via GET /jobs/{id}
     * Polls every 500ms, waiting for previous request to complete
     *
     * @param {number} jobId Job ID
     * @param {function} setProgressState Function to update progress
     * @param {function} onComplete Callback when job completes
     * @param {function} onProgressUpdate Callback on progress update
     * @param {object} statusMessages Custom status messages
     * @param {function} resetProgress Function to reset progress
     * @param {boolean} autoCloseNotifs Whether to auto-close notifications
     * @param {string} processType Process type for display
     */
    const pollJobStatus = useCallback(async ({
                                                 jobId,
                                                 setProgressState,
                                                 onComplete,
                                                 onProgressUpdate,
                                                 statusMessages,
                                                 resetProgress,
                                                 autoCloseNotifs = true,
                                                 processType
                                             }) => {
        // Initialize stop flag for this job
        if (!stopPollingRef.current.has(jobId)) {
            stopPollingRef.current.set(jobId, false);
        }
        if (!pollingInProgressRef.current.has(jobId)) {
            pollingInProgressRef.current.set(jobId, false);
        }

        // Prevent duplicate polling loops - if a polling loop is already active for this job, don't start another
        if (pollingLoopActiveRef.current.get(jobId)) {
            // Polling loop already running for this job, return
            return;
        }

        // Mark polling loop as active
        pollingLoopActiveRef.current.set(jobId, true);

        // Clear processed errors for this job
        processedFailedReasons.current.clear();

        while (!stopPollingRef.current.get(jobId)) {
            // Wait for previous request to complete
            while (pollingInProgressRef.current.get(jobId)) {
                await delay(100);
                if (stopPollingRef.current.get(jobId)) {
                    break;
                }
            }

            if (stopPollingRef.current.get(jobId)) {
                break;
            }

            try {
                pollingInProgressRef.current.set(jobId, true);
                let job;
                try {
                    job = await GetJob(jobId);
                } catch (e) {
                    // Job might have been removed (completed)
                    pollingInProgressRef.current.set(jobId, false);
                    job = null; // Treat error as job removal
                }
                pollingInProgressRef.current.set(jobId, false);

                // Check if job is completed (returns {completed: true, message: "..."})
                if (job && job.completed === true) {
                    // Job completed and was removed - handle completion
                    stopProcessingRef.current.set(jobId, true);
                    stopPollingRef.current.set(jobId, true);

                    // Remove from active jobs since job is completed and removed
                    setActiveJobs(prevJobs => prevJobs.filter(j => j.id !== jobId));

                    // Get job type BEFORE clearing processState (needed for progress reset)
                    const storedJobType = jobTypesRef.current.get(jobId);
                    const activeJob = activeJobs.find(j => j.id === jobId);
                    // Use processStateRef for reliable access to current state
                    const currentProcessState = processStateRef.current;
                    const jobType = storedJobType || activeJob?.type || currentProcessState?.type;

                    // Reset progress immediately (before clearing processState)
                    const progressSetter = getProgressSetter(jobType);
                    progressSetter(false);
                    if (typeof setProgressState === 'function') {
                        setProgressState(false);
                    }

                    // Restore document title
                    restoreDocumentTitle();

                    // Clear process state to reset button (no loading/pause state)
                    // Use processStateRef to check current state reliably
                    // Clear if jobId matches, or if type matches (safety check for edge cases)
                    if (currentProcessState?.jobId === jobId || currentProcessState?.type === jobType) {
                        setProcessState(null);
                    }

                    // Get stored statusMessages for this job
                    const storedStatusMessages = statusMessagesRef.current.get(jobId) || statusMessages;
                    // Show completion notification (only if not already shown and statusMessages.complete is not empty string)
                    const shouldShowNotification = !isManualActionRef.current && !completionNotificationShownRef.current.get(jobId) && storedStatusMessages?.complete !== '';
                    if (shouldShowNotification) {
                        const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                        const successMessage = DATASET_PROCESSES_INFO[jobType]?.success || storedStatusMessages?.complete || sprintf(__("%s process is complete.", 'limb-chatbot'), DATASET_PROCESSES_INFO[jobType]?.processName || jobType);

                        completionNotificationShownRef.current.set(jobId, true);
                        // Clean up stored job type and statusMessages
                        jobTypesRef.current.delete(jobId);
                        statusMessagesRef.current.delete(jobId);
                        handleSuccess(notifications.set, {
                            title: successMessage,
                            autoClose: autoCloseNotifs,
                        });
                    }

                    // Call resetProgress if registered (prefer stored callback over passed callback)
                    const storedResetProgress = resetProgressCallbacksRef.current.get(jobId);
                    if (typeof storedResetProgress === 'function') {
                        resetProgressCallbacksRef.current.delete(jobId);
                        setTimeout(() => {
                            storedResetProgress();
                        }, 200);
                    } else if (typeof resetProgress === 'function') {
                        // Fallback to passed callback if not stored
                        setTimeout(() => {
                            resetProgress();
                        }, 200);
                    }

                    // Update datasets list when job completes
                    // Use stored callback (prefer stored callback over passed callback)
                    const storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (typeof storedOnProgressUpdate === 'function') {
                        onProgressUpdateCallbacksRef.current.delete(jobId);
                        await storedOnProgressUpdate();
                    } else if (typeof onProgressUpdate === 'function') {
                        // Fallback to passed callback if not stored
                        await onProgressUpdate();
                    }

                    // Call onComplete and clean up the stored callback (prefer stored callback over passed callback)
                    const storedOnComplete = onCompleteCallbacksRef.current.get(jobId);
                    if (typeof storedOnComplete === 'function') {
                        onCompleteCallbacksRef.current.delete(jobId);
                        await storedOnComplete();
                    } else if (typeof onComplete === 'function') {
                        // Fallback to passed callback if not stored
                        await onComplete();
                    }

                    // Clean up completed_parents_count tracking
                    previousCompletedParentsCountRef.current.delete(jobId);
                    // Clean up errors length tracking
                    previousErrorsLengthRef.current.delete(jobId);

                    break;
                }

                if (!job) {
                    // Job was removed - this case handles cleanup when GET fails or job is removed
                    // Note: Completion callbacks are handled by other paths (job.completed or status check)
                    // so we don't call them here to avoid duplicate calls
                    break;
                }

                const status = job.status;
                const progress = job.progress_percent || 0;
                const errors = job.errors || [];
                const jobType = processType || job.type;

                // Get completed_parents_count from stats
                const completedParentsCount = job.stats?.completed_parents_count ?? null;

                // Get errors length
                const errorsLength = Array.isArray(errors) ? errors.length : 0;

                // Get the correct progress setter based on job type
                const progressSetter = getProgressSetter(jobType);
                // Update progress directly (0-100%)
                progressSetter(progress);
                // Also call the passed setProgressState for backward compatibility
                if (typeof setProgressState === 'function') {
                    setProgressState(progress);
                }

                // Handle errors
                handleJobErrors(errors, jobType);

                // Update process state
                // Don't update if this job is currently paused (to prevent loading bar from showing)
                // Check if stopPollingRef is set (indicating we've paused/canceled) or if current state is paused
                const isCurrentlyPaused = stopPollingRef.current.get(jobId) || processState?.jobId === jobId && processState?.status === 'pause';
                const processStatus = JOB_TO_PROCESS_STATUS[status] || 'start';

                // Only update process state if not currently paused (unless the new status is also paused)
                if (!isCurrentlyPaused || processStatus === 'pause') {
                    setProcessState(prevState => ({
                        ...prevState,
                        type: jobType,
                        status: processStatus,
                        jobId: job.id,
                        datasetIds: job.config?.dataset_ids || null,
                        startedAt: job.started_at ? parseDateString(job.started_at) : null,
                    }));
                }

                // Update active jobs
                setActiveJobs(prevJobs => {
                    return prevJobs.map(j => j.id === job.id ? job : j);
                });

                // Handle status
                if (status === JOB_STATUS_PROCESSING || status === JOB_STATUS_GENERATING_TASKS || status === JOB_STATUS_PENDING) {
                    // Ongoing - continue polling
                    // Update document title with progress
                    const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                    const processTitle = DATASET_PROCESSES_INFO[jobType]?.inPageTitle || DATASET_PROCESSES_INFO[jobType]?.processName || jobType;
                    if (pageSlug !== 'onboarding') {
                        updateDocumentProgressTitle(progress, (clampedProgress, originalTitle) => {
                            return `${clampedProgress}% ${processTitle} - ${originalTitle}`;
                        }, {resetOnInvalid: false});
                    }

                    // Update datasets list when completed_parents_count changes or errors length changes
                    // This applies to dataset_generating, dataset_sync, and dataset_delete jobs
                    // Use stored callback (prefer stored callback over passed callback)
                    // If passed callback exists but not stored, store it for future use
                    let storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (!storedOnProgressUpdate && typeof onProgressUpdate === 'function') {
                        // Store the passed callback if it's not already stored (e.g., on resume after page refresh)
                        onProgressUpdateCallbacksRef.current.set(jobId, onProgressUpdate);
                        storedOnProgressUpdate = onProgressUpdate;
                    }
                    const onProgressUpdateCallback = storedOnProgressUpdate || onProgressUpdate;

                    if (typeof onProgressUpdateCallback === 'function') {
                        let shouldUpdate = false;

                        // Check if completed_parents_count has changed
                        const previousCompletedParentsCount = previousCompletedParentsCountRef.current.get(jobId);
                        if (completedParentsCount !== null) {
                            if (previousCompletedParentsCount === undefined || completedParentsCount !== previousCompletedParentsCount) {
                                // completed_parents_count has changed, update datasets list
                                previousCompletedParentsCountRef.current.set(jobId, completedParentsCount);
                                shouldUpdate = true;
                            }
                        }

                        // Check if errors length has changed
                        const previousErrorsLength = previousErrorsLengthRef.current.get(jobId);
                        if (previousErrorsLength === undefined || errorsLength !== previousErrorsLength) {
                            // errors length has changed, update datasets list
                            previousErrorsLengthRef.current.set(jobId, errorsLength);
                            shouldUpdate = true;
                        }

                        // For jobs that don't track completed_parents_count (e.g., sitemap_scrape),
                        // call onProgressUpdate when progress changes to ensure updates during job execution
                        if (completedParentsCount === null) {
                            const previousProgress = previousProgressRef.current.get(jobId);
                            if (previousProgress === undefined || progress !== previousProgress) {
                                previousProgressRef.current.set(jobId, progress);
                                shouldUpdate = true;
                            }
                        }

                        if (shouldUpdate) {
                            await onProgressUpdateCallback();
                        }
                    }
                    // Wait before sending next request
                    await delay(JOB_POLLING_INTERVAL);
                } else if (status === JOB_STATUS_COMPLETED) {
                    // Complete - show completion notification
                    // Update datasets list when job completes
                    // Use stored callback (prefer stored callback over passed callback)
                    // Note: If callback was already called and deleted in the job.completed === true handler above,
                    // it won't be found here, so we skip the update to avoid duplicate calls
                    const storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (typeof storedOnProgressUpdate === 'function') {
                        onProgressUpdateCallbacksRef.current.delete(jobId);
                        await storedOnProgressUpdate();
                    }
                    // Note: We don't use fallback here to avoid duplicate calls if callback was already called

                    // Clean up progress tracking
                    previousProgressRef.current.delete(jobId);
                    // Clean up completed_parents_count tracking
                    previousCompletedParentsCountRef.current.delete(jobId);
                    // Clean up errors length tracking
                    previousErrorsLengthRef.current.delete(jobId);

                    // Reset progress immediately (before clearing processState)
                    const progressSetter = getProgressSetter(jobType);
                    progressSetter(false);
                    if (typeof setProgressState === 'function') {
                        setProgressState(false);
                    }

                    // Restore document title
                    restoreDocumentTitle();

                    // Clear process state to reset button (no loading/pause state)
                    // Use processStateRef to check current state reliably
                    const currentProcessState = processStateRef.current;
                    // Clear if jobId matches, or if type matches (safety check for edge cases)
                    if (currentProcessState?.jobId === jobId || currentProcessState?.type === jobType) {
                        setProcessState(null);
                    }

                    // Get stored statusMessages for this job
                    const storedStatusMessages = statusMessagesRef.current.get(jobId) || statusMessages;
                    // Show completion notification (only if not already shown and statusMessages.complete is not empty string)
                    const shouldShowNotification = !isManualActionRef.current && !completionNotificationShownRef.current.get(jobId) && storedStatusMessages?.complete !== '';
                    if (shouldShowNotification) {
                        const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                        const successMessage = DATASET_PROCESSES_INFO[jobType]?.success || storedStatusMessages?.complete || sprintf(__("%s process is complete.", 'limb-chatbot'), DATASET_PROCESSES_INFO[jobType]?.processName || jobType);
                        completionNotificationShownRef.current.set(jobId, true);
                        // Clean up stored job type and statusMessages
                        jobTypesRef.current.delete(jobId);
                        statusMessagesRef.current.delete(jobId);
                        handleSuccess(notifications.set, {
                            title: successMessage,
                            autoClose: autoCloseNotifs,
                        });
                    }

                    stopPollingRef.current.set(jobId, true);

                    // Call resetProgress if registered (prefer stored callback over passed callback)
                    const storedResetProgress = resetProgressCallbacksRef.current.get(jobId);
                    if (typeof storedResetProgress === 'function') {
                        resetProgressCallbacksRef.current.delete(jobId);
                        setTimeout(() => {
                            storedResetProgress();
                        }, 200);
                    } else if (typeof resetProgress === 'function') {
                        // Fallback to passed callback if not stored
                        setTimeout(() => {
                            resetProgress();
                        }, 200);
                    }

                    // Call onComplete and clean up the stored callback (prefer stored callback over passed callback)
                    const storedOnComplete = onCompleteCallbacksRef.current.get(jobId);
                    if (typeof storedOnComplete === 'function') {
                        onCompleteCallbacksRef.current.delete(jobId);
                        await storedOnComplete();
                    } else if (typeof onComplete === 'function') {
                        // Fallback to passed callback if not stored
                        await onComplete();
                    }
                    break;
                } else if (status === JOB_STATUS_PAUSED) {
                    // Paused - update document title with current progress
                    const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                    const processTitle = DATASET_PROCESSES_INFO[jobType]?.inPageTitle || DATASET_PROCESSES_INFO[jobType]?.processName || jobType;
                    if (pageSlug !== 'onboarding') {
                        updateDocumentProgressTitle(progress, (clampedProgress, originalTitle) => {
                            return `${clampedProgress}% ${processTitle} - ${originalTitle}`;
                        }, {resetOnInvalid: false});
                    }

                    // Get the correct progress setter based on job type
                    const progressSetter = getProgressSetter(jobType);
                    // Keep the progress percentage visible when paused (don't reset it to false)
                    // The loading animation will stop because the process state will be set to 'pause'
                    // The percentage is also shown in document title via updateDocumentTitleWithProgress above
                    progressSetter(progress);
                    // Also call the passed setProgressState for backward compatibility
                    if (typeof setProgressState === 'function') {
                        setProgressState(progress);
                    }

                    // Update datasets list when paused
                    // Use stored callback (prefer stored callback over passed callback)
                    const storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (typeof storedOnProgressUpdate === 'function') {
                        await storedOnProgressUpdate();
                    } else if (typeof onProgressUpdate === 'function') {
                        // Fallback to passed callback if not stored
                        await onProgressUpdate();
                    }

                    // Clean up progress tracking
                    previousProgressRef.current.delete(jobId);
                    // Clean up completed_parents_count tracking
                    previousCompletedParentsCountRef.current.delete(jobId);
                    // Clean up errors length tracking
                    previousErrorsLengthRef.current.delete(jobId);

                    // Show notification based on whether pause was manual or automatic
                    if (isManualActionRef.current) {
                        // Manual pause - show success notification
                        handleSuccess(notifications.set, {
                            title: statusMessages?.pause || sprintf(__("%s process is paused.", 'limb-chatbot'), processName),
                            autoClose: autoCloseNotifs,
                        });
                    } else {
                        // Automatic pause - show warning notification
                        handleWarning(false, notifications.set, {
                            title: statusMessages?.pause || sprintf(__("%s process has been paused.", 'limb-chatbot'), processName),
                            autoClose: autoCloseNotifs,
                        });
                    }
                    stopPollingRef.current.set(jobId, true);
                    break;
                } else if (status === JOB_STATUS_FAILED || status === JOB_STATUS_CANCELED) {
                    // Failed or Canceled - update datasets list
                    // Use stored callback (prefer stored callback over passed callback)
                    const storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (typeof storedOnProgressUpdate === 'function') {
                        onProgressUpdateCallbacksRef.current.delete(jobId);
                        await storedOnProgressUpdate();
                    } else if (typeof onProgressUpdate === 'function') {
                        // Fallback to passed callback if not stored
                        await onProgressUpdate();
                    }

                    // Clean up progress tracking
                    previousProgressRef.current.delete(jobId);
                    // Clean up completed_parents_count tracking
                    previousCompletedParentsCountRef.current.delete(jobId);
                    // Clean up errors length tracking
                    previousErrorsLengthRef.current.delete(jobId);

                    // Show notification based on whether cancel was manual or automatic
                    // Only show for canceled status, not failed (failed is handled by handleJobErrors)
                    if (status === JOB_STATUS_CANCELED) {
                        const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                        const processName = DATASET_PROCESSES_INFO[jobType]?.processName || jobType;
                        if (isManualActionRef.current) {
                            // Manual cancel - show success notification
                            handleSuccess(notifications.set, {
                                title: statusMessages?.cancel || sprintf(__("%s process is cancelled.", 'limb-chatbot'), processName),
                                autoClose: autoCloseNotifs,
                            });
                        } else {
                            // Automatic cancel - show warning notification
                            handleWarning(false, notifications.set, {
                                title: statusMessages?.cancel || sprintf(__("%s process has been cancelled.", 'limb-chatbot'), processName),
                                autoClose: autoCloseNotifs,
                            });
                        }
                    }
                    stopPollingRef.current.set(jobId, true);
                    if (resetProgress) {
                        setTimeout(() => {
                            resetProgress();
                        }, 200);
                    }
                    restoreDocumentTitle();
                    break;
                }
            } catch (e) {
                pollingInProgressRef.current.set(jobId, false);
                handleError(e, notifications.set, {
                    title: __("Failed to check process state.", 'limb-chatbot'),
                    description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
                    autoClose: autoCloseNotifs,
                });
                // Continue polling on error (might be temporary)
            }
        }

        // When polling loop exits due to stopPollingRef being set (e.g., job completed from POST request),
        // make sure we update datasets list if callback still exists (wasn't already called)
        // This handles the case where processJobSequentially detected completion and stopped the GET loop
        // Check if job actually completed by making a final GET request
        if (stopPollingRef.current.get(jobId)) {
            const storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
            const completionShown = completionNotificationShownRef.current.get(jobId);

            // Check if job actually completed by making a final GET request
            // This is more reliable than checking completionNotificationShownRef
            let jobCompleted = false;
            try {
                const finalJob = await GetJob(jobId);
                if (finalJob && (finalJob.completed === true || finalJob.status === JOB_STATUS_COMPLETED)) {
                    jobCompleted = true;
                } else if (finalJob && (finalJob.status === JOB_STATUS_PAUSED || finalJob.status === JOB_STATUS_FAILED || finalJob.status === JOB_STATUS_CANCELED)) {
                    // Job is paused/failed/canceled, don't update
                } else if (!finalJob) {
                    // Job was removed, likely completed
                    jobCompleted = true;
                }
            } catch (e) {
                // If we can't check, use completionNotificationShownRef as fallback
                jobCompleted = completionShown;
            }

            // Only update if callback exists AND job was completed (not paused/canceled)
            // For pause, we don't show completion notification, so we skip the update here
            // For cancel/fail, we delete the callback, so it won't exist here
            if (typeof storedOnProgressUpdate === 'function' && jobCompleted) {
                // Callback still exists and job completed, meaning job completed from POST request before GET loop could handle it
                // Update datasets list before exiting
                onProgressUpdateCallbacksRef.current.delete(jobId);
                try {
                    await storedOnProgressUpdate();
                } catch (e) {
                    console.error('Error in onProgressUpdate callback when stopping polling:', e);
                }
            }
        }

        // Mark polling loop as inactive when it exits
        pollingLoopActiveRef.current.set(jobId, false);
    }, [notifications, processState, setProcessState, setActiveJobs, getProgressSetter, handleJobErrors, updateDocumentProgressTitle, pageSlug, restoreDocumentTitle, getDatasetProcessesInfo]);

    /**
     * Process job sequentially via POST /jobs/{id}/process
     * Sends requests one by one, waiting for each to complete
     *
     * @param {number} jobId Job ID
     */
    const processJobSequentially = useCallback(async (jobId) => {
        // Initialize stop flag for this job
        if (!stopProcessingRef.current.has(jobId)) {
            stopProcessingRef.current.set(jobId, false);
        }
        if (!processingInProgressRef.current.has(jobId)) {
            processingInProgressRef.current.set(jobId, false);
        }

        // Note: We don't need separate stats polling - pollJobStatus already handles both status and stats
        // by calling GetJob which returns the full job object with stats included

        let errorCount = 0;
        const MAX_RETRIES = 4;

        while (!stopProcessingRef.current.get(jobId)) {
            // Wait for previous request to complete
            while (processingInProgressRef.current.get(jobId)) {
                await delay(100);
                if (stopProcessingRef.current.get(jobId)) {
                    break;
                }
            }

            if (stopProcessingRef.current.get(jobId)) {
                break;
            }

            try {
                processingInProgressRef.current.set(jobId, true);
                const job = await ProcessJob(jobId);
                processingInProgressRef.current.set(jobId, false);

                // Reset error count on successful request
                errorCount = 0;

                if (!job) {
                    stopProcessingRef.current.set(jobId, true);
                    break;
                }

                // Check if job is completed (returns {completed: true, message: "..."})
                if (job.completed === true) {
                    // Job completed and was removed - handle completion
                    stopProcessingRef.current.set(jobId, true);
                    stopPollingRef.current.set(jobId, true);

                    // Remove from active jobs since job is completed and removed
                    setActiveJobs(prevJobs => prevJobs.filter(j => j.id !== jobId));

                    // Get job type BEFORE clearing processState (needed for progress reset)
                    const storedJobType = jobTypesRef.current.get(jobId);
                    const activeJob = activeJobs.find(j => j.id === jobId);
                    // Use processStateRef for reliable access to current state
                    const currentProcessState = processStateRef.current;
                    const jobType = storedJobType || activeJob?.type || job.type || currentProcessState?.type;

                    // Reset progress immediately (before clearing processState)
                    const progressSetter = getProgressSetter(jobType);
                    progressSetter(false);

                    // Restore document title
                    restoreDocumentTitle();

                    // Clear process state to reset button (no loading/pause state)
                    // Use processStateRef to check current state reliably
                    // Clear if jobId matches, or if type matches (safety check for edge cases)
                    if (currentProcessState?.jobId === jobId || currentProcessState?.type === jobType) {
                        setProcessState(null);
                    }

                    // Get stored statusMessages for this job
                    const storedStatusMessages = statusMessagesRef.current.get(jobId);
                    // Show completion notification (only if not already shown and statusMessages.complete is not empty string)
                    if (!isManualActionRef.current && !completionNotificationShownRef.current.get(jobId) && storedStatusMessages?.complete !== '') {
                        const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
                        const successMessage = DATASET_PROCESSES_INFO[jobType]?.success || storedStatusMessages?.complete || sprintf(__("%s process is complete.", 'limb-chatbot'), DATASET_PROCESSES_INFO[jobType]?.processName || jobType);

                        completionNotificationShownRef.current.set(jobId, true);
                        // Clean up stored job type and statusMessages
                        jobTypesRef.current.delete(jobId);
                        statusMessagesRef.current.delete(jobId);
                        handleSuccess(notifications.set, {
                            title: successMessage,
                            autoClose: true,
                        });
                    }

                    // Call resetProgress if registered
                    const resetProgress = resetProgressCallbacksRef.current.get(jobId);
                    if (typeof resetProgress === 'function') {
                        resetProgressCallbacksRef.current.delete(jobId);
                        setTimeout(() => {
                            resetProgress();
                        }, 200);
                    }

                    // IMPORTANT: Call onProgressUpdate BEFORE onComplete to ensure data is updated before cleanup
                    // This prevents data loss when onComplete deletes data that onProgressUpdate needs
                    const onProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
                    if (typeof onProgressUpdate === 'function') {
                        onProgressUpdateCallbacksRef.current.delete(jobId);
                        try {
                            await onProgressUpdate();
                        } catch (e) {
                            console.error('Error in onProgressUpdate callback:', e);
                        }
                    }

                    // Call onComplete callback if registered (after onProgressUpdate completes)
                    const onComplete = onCompleteCallbacksRef.current.get(jobId);
                    if (typeof onComplete === 'function') {
                        onCompleteCallbacksRef.current.delete(jobId);
                        try {
                            await onComplete();
                        } catch (e) {
                            console.error('Error in onComplete callback:', e);
                        }
                    }

                    // Clean up completed_parents_count tracking
                    previousCompletedParentsCountRef.current.delete(jobId);
                    // Clean up errors length tracking
                    previousErrorsLengthRef.current.delete(jobId);

                    break;
                }

                // If response doesn't have status, it might be invalid
                if (!job.status) {
                    stopProcessingRef.current.set(jobId, true);
                    break;
                }

                const status = job.status;

                // Update active jobs (only if job object has an id, meaning it wasn't removed)
                if (job.id) {
                    setActiveJobs(prevJobs => {
                        const existing = prevJobs.find(j => j.id === job.id);
                        if (existing) {
                            return prevJobs.map(j => j.id === job.id ? job : j);
                        } else {
                            return [...prevJobs, job];
                        }
                    });
                }

                // Stop if status indicates process is not running
                if (status === JOB_STATUS_PENDING ||
                    status === JOB_STATUS_PAUSED ||
                    status === JOB_STATUS_COMPLETED ||
                    status === JOB_STATUS_FAILED ||
                    status === JOB_STATUS_CANCELED) {
                    // For paused/failed/canceled, just stop processing
                    stopProcessingRef.current.set(jobId, true);
                    break;
                }

                // Continue processing if status is processing or generating_tasks
                if (status === JOB_STATUS_PROCESSING || status === JOB_STATUS_GENERATING_TASKS) {
                    // Continue loop
                } else {
                    // Unknown status, stop
                    stopProcessingRef.current.set(jobId, true);
                    break;
                }
            } catch (e) {
                processingInProgressRef.current.set(jobId, false);
                errorCount++;

                // Only stop and show notification after 4 failed attempts
                if (errorCount >= MAX_RETRIES) {
                    stopProcessingRef.current.set(jobId, true);
                    stopPollingRef.current.set(jobId, true);
                    restoreDocumentTitle();
                    handleError(e, notifications.set, {
                        title: __("Failed to process job.", 'limb-chatbot'),
                        description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
                    });
                    break;
                }
                // Otherwise, continue retrying (will retry on next loop iteration)
            }
        }
    }, [notifications, processState, setProcessState, setActiveJobs, activeJobs, restoreDocumentTitle, getDatasetProcessesInfo]);

    /**
     * Follow a job (start polling and processing)
     *
     * @param {object} options Options object
     */
    const followJob = useCallback(async ({
                                             jobId,
                                             setProgressState,
                                             onComplete,
                                             onProgressUpdate,
                                             statusMessages,
                                             resetProgress,
                                             autoCloseNotifs = true,
                                             processType
                                         }) => {
        // Reset manual action tracking
        isManualActionRef.current = false;

        // Reset completion notification tracking for this job
        completionNotificationShownRef.current.set(jobId, false);

        // Store job type so we can retrieve it when job completes (even if job is removed)
        if (processType) {
            jobTypesRef.current.set(jobId, processType);
        }

        // Store onComplete, resetProgress, and onProgressUpdate callbacks so they can be called when job completes
        if (typeof onComplete === 'function') {
            onCompleteCallbacksRef.current.set(jobId, onComplete);
        }
        if (typeof resetProgress === 'function') {
            resetProgressCallbacksRef.current.set(jobId, resetProgress);
        }
        if (typeof onProgressUpdate === 'function') {
            onProgressUpdateCallbacksRef.current.set(jobId, onProgressUpdate);
        }
        // Store statusMessages so it's available when job completes
        if (statusMessages) {
            statusMessagesRef.current.set(jobId, statusMessages);
        }

        // Show progress immediately when job starts (even if 0%)
        setProgressState(0);
        updateDocumentTitleWithProgress(0, processType);

        // Start sequential processing
        processJobSequentially(jobId);

        // Start polling
        await pollJobStatus({
            jobId,
            setProgressState,
            onComplete,
            onProgressUpdate,
            statusMessages,
            resetProgress,
            autoCloseNotifs,
            processType
        });
    }, [processJobSequentially, pollJobStatus, updateDocumentTitleWithProgress]);

    /**
     * Create a job
     *
     * @param {object} jobData Job data
     * @return {Promise<object>} Created job
     */
    const createJob = useCallback(async (jobData) => {
        try {
            const job = await CreateJob(jobData);

            // Validate that job was created successfully
            if (!job || !job.id) {
                const error = new Error(__("Job creation failed: Invalid response from server.", 'limb-chatbot'));
                handleError(error, notifications.set, {
                    title: __("Failed to create job.", 'limb-chatbot'),
                    description: __("The server did not return a valid job. Please check and try again.", 'limb-chatbot'),
                });
                return null;
            }

            // Add to active jobs
            setActiveJobs(prevJobs => {
                const existing = prevJobs.find(j => j.id === job.id);
                if (existing) {
                    return prevJobs.map(j => j.id === job.id ? job : j);
                } else {
                    return [...prevJobs, job];
                }
            });

            return job;
        } catch (e) {
            handleError(e, notifications.set, {
                title: __("Failed to create job.", 'limb-chatbot'),
                description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
            });
            throw e;
        }
    }, [notifications, setActiveJobs]);

    /**
     * Pause the current process
     */
    const pauseProcess = useCallback(() => {
        const jobId = processState?.jobId;
        if (!jobId) {
            return;
        }

        setPausing(true);
        PauseJob(jobId).then(job => {
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            // Get process name
            const processName = DATASET_PROCESSES_INFO[processState.type]?.processName || processState.type;

            // Get progress to show
            const progressToShow = job.progress_percent || 0;

            // Update document title with current progress
            updateDocumentTitleWithProgress(progressToShow, processState.type);

            // Get appropriate progress setter based on job type
            const progressSetter = getProgressSetter(processState.type);
            // Keep the progress percentage visible when paused (don't reset it to false)
            // The loading animation will stop because the process state is 'pause'
            // The percentage is also shown in document title via updateDocumentTitleWithProgress above
            progressSetter(progressToShow);

            // Update process state
            setProcessState(prevState => ({
                ...prevState,
                status: 'pause',
            }));

            // Update active jobs
            setActiveJobs(prevJobs => prevJobs.map(j => j.id === job.id ? job : j));

            // Stop polling and processing
            stopPollingRef.current.set(jobId, true);
            stopProcessingRef.current.set(jobId, true);

            // Mark as manual action to avoid duplicate notifications
            isManualActionRef.current = true;

            // Show success notification for manual pause
            handleSuccess(notifications.set, {
                title: sprintf(__("%s process is paused.", 'limb-chatbot'), processName),
            });

            setPausing(false);
        }).catch((e) => {
            setPausing(false);
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            handleError(e, notifications.set, {
                title: sprintf(__("Failed to pause %s.", 'limb-chatbot'), DATASET_PROCESSES_INFO[processState.type]?.processName?.toLowerCase() || processState.type),
                description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
            });
        });
    }, [processState, notifications, setPausing, setProcessState, setActiveJobs, updateDocumentTitleWithProgress, getProgressSetter, getDatasetProcessesInfo]);

    /**
     * Resume the current process
     */
    const resumeProcess = useCallback(async () => {
        // Read current processState from ref for reliable access in async context
        // This ensures we always get the latest state even if the component re-renders during async operations
        const currentState = processStateRef.current || processState;
        const jobId = currentState?.jobId;

        if (!jobId) {
            return;
        }

        if (currentState?.status !== 'pause') {
            // If status is not 'pause', the job might have already been resumed or is in a different state
            // This is not an error, just return silently
            return;
        }

        // Prevent duplicate resume calls - use a flag to track if resume is in progress
        // Initialize the ref if it doesn't exist
        if (!resumeInProgressRef.current.has(jobId)) {
            resumeInProgressRef.current.set(jobId, false);
        }

        // Check if resume is already in progress
        const isResumeInProgress = resumeInProgressRef.current.get(jobId);
        if (isResumeInProgress) {
            return;
        }

        // Set resume in progress flag
        resumeInProgressRef.current.set(jobId, true);

        // Store processType from current state before async operations
        const processType = currentState?.type;

        // Use a ref to track if we've cleared the resume flag to avoid double-clearing
        const resumeClearedRef = {current: false};
        const clearResumeFlag = () => {
            if (!resumeClearedRef.current) {
                resumeClearedRef.current = true;
                resumeInProgressRef.current.set(jobId, false);
            }
        };

        // Set pausing state to show loading indicator and disable button during resume
        // IMPORTANT: This must be reset in ALL code paths (success, error, early returns)
        setPausing(true);
        try {
            // Ensure all previous loops have stopped by setting stop flags first
            stopPollingRef.current.set(jobId, true);
            stopProcessingRef.current.set(jobId, true);

            // Wait for any running loops to actually stop
            let waitCount = 0;
            while ((pollingInProgressRef.current.get(jobId) ||
                processingInProgressRef.current.get(jobId)) &&
            waitCount < 50) { // Max 5 seconds wait
                await delay(100);
                waitCount++;
            }

            const job = await ResumeJob(jobId);

            // Verify the job was actually resumed (status should not be 'paused')
            if (job.status === JOB_STATUS_PAUSED) {
                // Job is still paused, something went wrong
                clearResumeFlag();
                setPausing(false);
                handleError(false, notifications.set, {
                    title: __("Failed to resume job.", 'limb-chatbot'),
                    description: __("The job is still paused. Please try again.", 'limb-chatbot'),
                });
                return;
            }

            // Update active jobs
            setActiveJobs(prevJobs => prevJobs.map(j => j.id === job.id ? job : j));

            // Get job type from job object or fallback to stored processType
            const jobType = job.type || processType;

            // Update process state
            setProcessState(prevState => ({
                ...prevState,
                type: jobType,
                status: 'start',
                jobId: job.id,
            }));

            // Reset stop flags to allow new loops to start
            stopPollingRef.current.set(jobId, false);
            stopProcessingRef.current.set(jobId, false);

            // Reset in-progress flags to ensure clean start
            pollingInProgressRef.current.set(jobId, false);
            processingInProgressRef.current.set(jobId, false);

            // Reset polling loop active flags to allow new loops to start
            pollingLoopActiveRef.current.set(jobId, false);

            // Get appropriate progress setter based on job type
            const progressSetter = getProgressSetter(jobType);

            // Get progress to show
            const progressToShow = job.progress_percent || 0;

            // Get process name for notifications and document title
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            const processName = DATASET_PROCESSES_INFO[jobType]?.processName || jobType;

            // Show progress immediately when resuming
            progressSetter(progressToShow);
            if (pageSlug !== 'onboarding') {
                const processTitle = DATASET_PROCESSES_INFO[jobType]?.inPageTitle || processName;
                updateDocumentProgressTitle(progressToShow, (clampedProgress, originalTitle) => {
                    return `${clampedProgress}% ${processTitle} - ${originalTitle}`;
                }, {resetOnInvalid: false});
            }

            // Resume processing and polling
            processJobSequentially(jobId);

            // Get the stored onProgressUpdate callback if it exists (for datasets list updates)
            // Try both jobId (from processState) and job.id (from resumed job) in case they differ
            // The callback should still be in the ref from when followJob was first called
            // If it's not there (e.g., page refresh), we'll need to get it from the component
            let storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(job.id);
            if (!storedOnProgressUpdate && jobId !== job.id) {
                // Try the original jobId if it's different from job.id
                storedOnProgressUpdate = onProgressUpdateCallbacksRef.current.get(jobId);
            }

            // Re-store the callback in the ref if it exists (in case it was lost, e.g., on page refresh)
            // This ensures it's available during polling and for completion updates
            // Use job.id (from resumed job) as the key
            if (typeof storedOnProgressUpdate === 'function') {
                onProgressUpdateCallbacksRef.current.set(job.id, storedOnProgressUpdate);
            }

            // Restart polling - pass the stored callback so it can be used for updates
            // pollJobStatus will use the stored callback from the ref (which we just ensured is set)
            // If callback is passed but not stored, pollJobStatus will store it automatically
            pollJobStatus({
                jobId: job.id,
                setProgressState: progressSetter,
                onComplete: () => {
                    setProcessState(null);
                    // Clear resume in progress flag when job completes
                    clearResumeFlag();
                },
                onProgressUpdate: storedOnProgressUpdate || (() => {
                }),
                statusMessages: {
                    complete: DATASET_PROCESSES_INFO[jobType]?.success || sprintf(__("%s process is complete.", 'limb-chatbot'), processName),
                    pause: sprintf(__("%s process has been paused.", 'limb-chatbot'), processName),
                    cancel: sprintf(__("%s process has been cancelled.", 'limb-chatbot'), processName)
                },
                resetProgress: () => {
                },
                autoCloseNotifs: true,
                processType: jobType,
            });

            // Clear resume flag immediately after successfully starting the job
            // This allows the button to be clickable again if needed
            // The flag was only needed to prevent duplicate calls during the resume operation
            clearResumeFlag();
            setPausing(false);
        } catch (e) {
            // Clear resume in progress flag on error
            clearResumeFlag();
            setPausing(false);
            const errorJobType = processType || 'job';
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            handleError(e, notifications.set, {
                title: sprintf(__("Failed to resume %s.", 'limb-chatbot'), DATASET_PROCESSES_INFO[errorJobType]?.processName?.toLowerCase() || errorJobType),
                description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
            });
        }
    }, [processStateRef, processState, notifications, setPausing, setProcessState, setActiveJobs, updateDocumentTitleWithProgress, getProgressSetter, processJobSequentially, pollJobStatus, getDatasetProcessesInfo]);

    /**
     * Cancel the current process
     */
    const cancelProcess = useCallback(() => {
        const jobId = processState?.jobId;
        if (!jobId) {
            return;
        }

        setCanceling(true);

        // Cancel the main job
        CancelJob(jobId).then(res => {
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            const processName = DATASET_PROCESSES_INFO[processState.type]?.processName || processState.type;

            // Process canceled manually - show success notification
            handleSuccess(notifications.set, {
                title: sprintf(__("%s process is cancelled.", 'limb-chatbot'), processName),
            });

            // Mark as manual action to avoid duplicate notifications
            isManualActionRef.current = true;

            // Stop polling and processing
            stopPollingRef.current.set(jobId, true);
            stopProcessingRef.current.set(jobId, true);

            // Remove from active jobs
            setActiveJobs(prevJobs => prevJobs.filter(j => j.id !== jobId));

            // Reset process state
            setProcessState(null);
            restoreDocumentTitle();
            setCanceling(false);
            setSaving(false);
        }).catch((e) => {
            setCanceling(false);
            const DATASET_PROCESSES_INFO = getDatasetProcessesInfo();
            handleError(e, notifications.set, {
                title: sprintf(__("Failed to cancel %s.", 'limb-chatbot'), DATASET_PROCESSES_INFO[processState.type]?.processName?.toLowerCase() || processState.type),
                description: e.message ? __(e.message, 'limb-chatbot') : __("Please check and try again.", 'limb-chatbot'),
            });
        });
    }, [processState, notifications, setCanceling, setSaving, setProcessState, setActiveJobs, restoreDocumentTitle, getDatasetProcessesInfo]);

    // Memoize context value to prevent unnecessary re-renders
    const contextValue = useMemo(() => ({
        // Process state
        processState,
        activeJobs,

        // Progress states
        generatingProgress,
        syncingProgress,
        deletingProgress,

        // Loading states
        saving,
        pausing,
        canceling,
        preparing,
        syncing,
        deleting,

        // Process management functions
        checkActiveProcess,
        pauseProcess,
        resumeProcess,
        cancelProcess,
        createJob,
        followJob,
        processJobSequentially,
        handleJobErrors,

        // Progress update functions
        setGeneratingProgress,
        setSyncingProgress,
        setDeletingProgress,

        // Loading state functions
        setSaving,
        setPausing,
        setCanceling,
        setPreparing,
        setSyncing,
        setDeleting,

        // Internal state setters
        setProcessState,
        setActiveJobs,
    }), [
        // State values
        processState,
        activeJobs,
        generatingProgress,
        syncingProgress,
        deletingProgress,
        saving,
        pausing,
        canceling,
        preparing,
        syncing,
        deleting,
        // Functions
        checkActiveProcess,
        pauseProcess,
        resumeProcess,
        cancelProcess,
        createJob,
        followJob,
        processJobSequentially,
        handleJobErrors,
        // Setters (these are stable from useState)
        setGeneratingProgress,
        setSyncingProgress,
        setDeletingProgress,
        setSaving,
        setPausing,
        setCanceling,
        setPreparing,
        setSyncing,
        setDeleting,
        setProcessState,
        setActiveJobs,
    ]);

    return (
        <BackgroundProcessContext.Provider value={contextValue}>
            {children}
        </BackgroundProcessContext.Provider>
    );
}

// Hook for using the context
export function useBackgroundProcess() {
    const context = useContext(BackgroundProcessContext);
    if (!context) {
        throw new Error('useBackgroundProcess must be used within a BackgroundProcessProvider');
    }
    return context;
}

export default BackgroundProcessContext;
