function batchResults = runBatch(workerFcn, jobs, options)
%RUNBATCH Execute a list of jobs in serial or parallel.
%   batchResults = runBatch(workerFcn, jobs) executes each jobs(i).args cell
%   with workerFcn and returns a cell array aligned with the input jobs.
%
%   Supported job fields:
%       .args   cell array of positional arguments for workerFcn
%       .label  string/char used for progress messages
%       .meta   arbitrary metadata for wrapper-specific bookkeeping

    arguments
        workerFcn (1,1) function_handle
        jobs (1,:) struct = struct('args', {})
        options.mode = processingMode.serial
        options.waitbar (1,1) logical = true
        options.waitbarMessage (1,1) string = "Processing Jobs..."
        options.numWorkers (1,1) double {mustBeNonnegative, mustBeInteger} = 0
        options.idleTimeout (1,1) double {mustBePositive} = 300
        options.cancelExistingQueue (1,1) logical = true
        options.resultHandler = []
        options.errorHandler = []
    end

    mode = normalizeProcessingMode(options.mode);
    jobs = normalizeJobs(jobs);
    nJobs = numel(jobs);
    batchResults = cell(1, nJobs);
    h = [];

    if nJobs == 0
        return
    end

    if options.waitbar
        h = waitbar(0, char(options.waitbarMessage));
        cleanupWaitbar = onCleanup(@() closeWaitbar(h));
    end

    switch mode
        case processingMode.parallel
            pool = setupParallelPool(options.numWorkers, options.idleTimeout, options.cancelExistingQueue);
            futures = parallel.FevalFuture.empty(nJobs, 0);

            for jobIdx = 1:nJobs
                fprintf('[%s] Submitted.\n', jobs(jobIdx).label);
                futures(jobIdx) = parfeval(pool, workerFcn, 1, jobs(jobIdx).args{:});
            end

            consumedIdx = false(nJobs, 1);
            for completedCount = 1:nJobs
                try
                    [jobIdx, jobResult] = fetchNext(futures);
                    consumedIdx(jobIdx) = true;
                    batchResults{jobIdx} = jobResult;
                    fprintf('[%s] Completed (%d/%d).\n', jobs(jobIdx).label, completedCount, nJobs);

                    if ~isempty(options.resultHandler)
                        options.resultHandler(jobResult, jobs(jobIdx), jobIdx);
                    end
                catch fetchErr
                    errorJobIdx = findErroredFuture(futures, consumedIdx);
                    if isempty(errorJobIdx)
                        rethrow(fetchErr);
                    end

                    consumedIdx(errorJobIdx) = true;
                    errInfo = extractFutureError(futures(errorJobIdx));
                    batchResults{errorJobIdx} = errInfo;

                    if ~isempty(options.errorHandler)
                        options.errorHandler(errInfo, jobs(errorJobIdx), errorJobIdx);
                    else
                        defaultErrorHandler(errInfo, jobs(errorJobIdx), errorJobIdx);
                    end
                end

                updateWaitbar(options.waitbar, h, completedCount, nJobs);
            end

        case processingMode.serial
            for jobIdx = 1:nJobs
                try
                    fprintf('[%s] Running.\n', jobs(jobIdx).label);
                    jobResult = workerFcn(jobs(jobIdx).args{:});
                    batchResults{jobIdx} = jobResult;
                    fprintf('[%s] Completed (%d/%d).\n', jobs(jobIdx).label, jobIdx, nJobs);

                    if ~isempty(options.resultHandler)
                        options.resultHandler(jobResult, jobs(jobIdx), jobIdx);
                    end
                catch ME
                    batchResults{jobIdx} = ME;

                    if ~isempty(options.errorHandler)
                        options.errorHandler(ME, jobs(jobIdx), jobIdx);
                    else
                        defaultErrorHandler(ME, jobs(jobIdx), jobIdx);
                    end
                end

                updateWaitbar(options.waitbar, h, jobIdx, nJobs);
            end

        otherwise
            error('runBatch:UnknownMode', 'Unknown processing mode "%s".', string(mode));
    end
end

function mode = normalizeProcessingMode(mode)
    if isa(mode, 'processingMode')
        return
    end

    modeString = string(mode);
    switch modeString
        case "serial"
            mode = processingMode.serial;
        case "parallel"
            mode = processingMode.parallel;
        otherwise
            error('runBatch:InvalidMode', 'Unsupported processing mode "%s".', modeString);
    end
end

function jobs = normalizeJobs(jobs)
    if isempty(jobs)
        jobs = struct('args', {}, 'label', {}, 'meta', {});
        return
    end

    for jobIdx = 1:numel(jobs)
        if ~isfield(jobs, 'args') || isempty(jobs(jobIdx).args)
            jobs(jobIdx).args = {};
        end

        if ~iscell(jobs(jobIdx).args)
            error('runBatch:InvalidArgs', 'jobs(%d).args must be a cell array.', jobIdx);
        end

        if ~isfield(jobs, 'label') || isempty(jobs(jobIdx).label)
            jobs(jobIdx).label = sprintf('Job %d', jobIdx);
        else
            jobs(jobIdx).label = string(jobs(jobIdx).label);
        end

        if ~isfield(jobs, 'meta')
            jobs(jobIdx).meta = struct();
        end
    end
end

function updateWaitbar(useWaitbar, h, completedCount, totalCount)
    if ~useWaitbar || ~isgraphics(h)
        return
    end

    waitbar(completedCount / totalCount, h, ...
        sprintf('Completed %d/%d jobs', completedCount, totalCount));
    drawnow;
end

function closeWaitbar(h)
    if isgraphics(h)
        delete(h);
    end
end

function pool = setupParallelPool(numWorkers, idleTimeout, cancelExistingQueue)
    pool = gcp('nocreate');

    if isempty(pool)
        if numWorkers > 0
            pool = parpool('local', numWorkers, 'IdleTimeout', idleTimeout);
        else
            pool = parpool('local', 'IdleTimeout', idleTimeout);
        end
    elseif numWorkers > 0 && pool.NumWorkers ~= numWorkers
        delete(pool);
        pool = parpool('local', numWorkers, 'IdleTimeout', idleTimeout);
    end

    if cancelExistingQueue
        queue = pool.FevalQueue;
        if ~isempty(queue.QueuedFutures) || ~isempty(queue.RunningFutures)
            queuedCount = numel(queue.QueuedFutures) + numel(queue.RunningFutures);
            cancelAll(queue);
            fprintf('Canceled %d unfinished jobs from the current pool queue.\n', queuedCount);
        end
    end
end

function errorJobIdx = findErroredFuture(futures, consumedIdx)
    readMask = arrayfun(@(future) future.Read, futures).';
    errorJobIdx = find(readMask & ~consumedIdx, 1);
end

function errInfo = extractFutureError(future)
    errInfo = future.Error;
    if iscell(errInfo)
        errInfo = errInfo{1};
    end
end

function defaultErrorHandler(ME, job, jobIndex)
    fprintf('[%s | #%d] ERROR [%s]: %s\n', job.label, jobIndex, ME.identifier, ME.message);
    for st = ME.stack'
        fprintf('    %s:%d (%s)\n', st.file, st.line, st.name);
    end
end