function [OK, out] = remoteBlockExecutionThreads(action, varargin)
%remoteBlockExecutionThreads controls the SPMD block execution on the
%workers of a ThreadPool.

% Copyright 2021-2022 The MathWorks, Inc.

    persistent EXECUTION_HANDLES PRELUDE_EXCEPTION

    mlock;

    out = [];

    dctSchedulerMessage(5, 'Entering remoteBlockExecutionThreads with action "%s".', action);
    logOnExit = onCleanup(@() dctSchedulerMessage(5, 'Leaving remoteBlockExecutionThreads with action "%s".', action));

    try
        switch action
          case 'prelude'
            try
                % Switch to use threads message passing functions.
                parallel.internal.mpi.registerFunctions('threadsmpi');
                preEvalException = [];
            catch preEvalException
            end

            % Set up defaults in case iPrelude errors out
            EXECUTION_HANDLES = [];
            PRELUDE_EXCEPTION = [];

            try
                [EXECUTION_HANDLES, PRELUDE_EXCEPTION] = iPrelude(varargin{:}, preEvalException);
            catch E
                % Unexpected error in prelude. We have no way to continue, so abort.
                dctSchedulerMessage(0, 'Unexpected error during SPMD prelude: %s', E.identifier)
            end

          case 'interruptibleExecution'
            % In this case, we expect interruption, and so do not rely on modifying any
            % of the persistent data.
            if isempty(PRELUDE_EXCEPTION)
                bodyFcn = EXECUTION_HANDLES{1};

                bodyFcn();

                iSetFinished();
            else

                iSetFinished();

                % Don't even try to execute the block
                throw(PRELUDE_EXCEPTION);
            end
          case 'postlude'
            iSetFinished();
            iSetRunning();

            % On cleanup, reset the MPI functions to their defaults.
            cleanupObj = onCleanup(@() parallel.internal.mpi.registerFunctions('mpi_mi'));

            if isempty(PRELUDE_EXCEPTION)
                getOutFcn = EXECUTION_HANDLES{2};
                out = iPackOutputs(getOutFcn());
            else
                % No block execution - bad setup
                out = [];
            end
            % Done with EXECUTION_HANDLES - clear it now because it can cause
            % data and Composites to stay in memory (g2420794).
            EXECUTION_HANDLES = [];

            % Reset parfor_depth unconditionally.
            parfor_depth(0);

        end
        OK = true;
    catch E
        OK = false;
        % This is where we are actually evaluating the user code, so the
        % error that we send back needs to include the stack info
        % so we can stitch together a new exception later.
        stackFilesToIgnore = {'LocalSpmdExecutor', 'remoteBlockExecutionThreads'};
        out = ParallelException.hBuildFromRemoteException(E, stackFilesToIgnore);
    end
end

% iPrelude - unpack inputs, setup communicator and clear any applicable
% keys.
function [fcnHandles, preludeExcept] = iPrelude(commId, fcnHandles, keysToClear, preludeExcept)

% Clear any keys that are no longer in use.
    pClearIfNecessary(keysToClear);

    if ~isempty(preludeExcept)
        return;
    end

    % Set the communicator for this worker.
    try
        parallel.internal.mpi.setThreadCommunicator(commId);
    catch E
        preludeExcept = E;
    end

    if ~isempty(preludeExcept)
        return;
    end

    % Unpack the inputs for later
    try
        unpackInFcn = fcnHandles{3};
        unpackInFcn(@spmdlang.AbstractSpmdExecutor.unpack);
    catch E
        if strcmp(E.identifier, 'MATLAB:UndefinedFunction')
            preludeExcept = MException(message('parallel:lang:spmd:SourceCodeNotAvailableOnWorker', spmdIndex));
        else
            rethrow(E);
        end
    end

end

% iPackOutputs - pack up and return the body outputs
function outCellParcel = iPackOutputs(outCell)
    outCellParcel = cell(1, length(outCell));
    for ii=1:length(outCell)
        try
            if isempty(outCell{ii})
                % Do nothing - cell already empty
            else
                data = outCell{ii}{1};
                key = spmdlang.ValueStore.store(data);
                [fcn, data] = getRemoteFromSPMD(data);
                outCellParcel{ii} = spmdlang.ReturnableParcel(key, fcn, data);
            end
        catch E
            warning(message('parallel:lang:spmd:SpmdOutputs', getReport(E)));
        end
    end
end

function iSetFinished()

    parallel.internal.mpi.setThreadCommunicatorStatus("finished");

end

function iSetRunning()
    % Note, this is a collective operation.
    parallel.internal.mpi.setThreadCommunicatorStatus("running");

end
