%%% @doc
%%% Conveniences for spawning and awaiting tasks.
%%%
%%% Tasks are processes meant to execute one particular
%%% action throughout their lifetime, often with little or no
%%% communication with other processes. The most common use case
%%% for tasks is to convert sequential code into concurrent code
%%% by computing a value asynchronously:
%%%
%%% ```
%%% Task = task:async(fun() -> do_some_work() end),
%%% Res = do_some_other_work(),
%%% Res + task:await(Task).
%%% '''
%%%
%%% Tasks spawned with async/1 can be awaited on by their caller
%%% process (and only their caller) as shown in the example above.
%%% They are implemented by spawning a process that sends a message
%%% to the caller once the given computation is performed.
%%%
%%% Compared to plain processes, started with spawn/1, tasks
%%% include monitoring metadata and logging in case of errors.
%%%
%%% Besides async/1 and await/2, tasks can also be
%%% started as part of a supervision tree and dynamically spawned
%%% on remote nodes.
%%%
%%% == async and await ==
%%%
%%% One of the common uses of tasks is to convert sequential code
%%% into concurrent code with async/1 while keeping its semantics.
%%% When invoked, a new process will be created, linked and monitored
%%% by the caller. Once the task action finishes, a message will be sent
%%% to the caller with the result.
%%%
%%% await/2 is used to read the message sent by the task.
%%%
%%% There are two important things to consider when using async:
%%%
%%% 1. If you are using async tasks, you **must await** a reply
%%%    as they are *always* sent. If you are not expecting a reply,
%%%    consider using start_link/1 as detailed below.
%%%
%%% 2. Async tasks link the caller and the spawned process. This
%%%    means that, if the caller crashes, the task will crash
%%%    too and vice-versa. This is on purpose: if the process
%%%    meant to receive the result no longer exists, there is
%%%    no purpose in completing the computation.
%%%
%%% == Tasks are processes ==
%%%
%%% Tasks are processes and so data will need to be completely copied
%%% to them. Consider extracting only the necessary data before creating
%%% the task, or move the data loading altogether to the task.
%%%
%%% @end
-module(task).

%% API
-export([async/1, async/3]).
-export([await/1, await/2]).
-export([yield/1, yield/2]).
-export([start_link/1, start_link/3]).
-export([start/1, start/3]).
-export([shutdown/1, shutdown/2]).
-export([completed/1]).

%% Types
-record(task, {
    pid :: pid() | undefined,
    ref :: reference(),
    owner :: pid(),
    mfa :: mfa()
}).

-opaque task() :: #task{}.
-type task_fun() :: fun(() -> term()).

-export_type([task/0]).

%%% @doc
%%% Starts a task that must be awaited on.
%%%
%%% This function spawns a new process linked to the caller
%%% that will execute the given function. A task record is returned
%%% containing the information needed to await the result.
%%%
%%% The spawned process will send a message to the caller once the
%%% computation is finished.
%%%
%%% This function should be used when you want to run a computation
%%% concurrently and retrieve its result. The caller process must
%%% call await/1 or await/2 to retrieve the result.
%%%
%%% == Example ==
%%%
%%% ```
%%% Task = task:async(fun() -> 1 + 1 end),
%%% 2 = task:await(Task).
%%% '''
%%% @end
-spec async(task_fun()) -> task().
async(Fun) when is_function(Fun, 0) ->
    async(erlang, apply, [Fun, []]).

%%% @doc
%%% Starts a task that must be awaited on with the given module, function, and arguments.
%%%
%%% Similar to async/1 except the function to be started is
%%% specified by the given Module, Function, and Args.
%%% The Module, Function, and its arity are stored as
%%% a tuple in the mfa field for reflection purposes.
%%% @end
-spec async(module(), atom(), [term()]) -> task().
async(Module, Function, Args) when is_atom(Module), is_atom(Function), is_list(Args) ->
    Owner = self(),
    Ref = make_ref(),

    Pid = spawn_link(fun() ->
        put('$initial_call', {Module, Function, length(Args)}),
        put('$ancestors', [Owner]),
        try
            Result = apply(Module, Function, Args),
            Owner ! {Ref, Result}
        catch
            Class:Reason:Stacktrace ->
                Owner ! {Ref, {Class, Reason, Stacktrace}}
        end
    end),

    #task{pid = Pid, ref = Ref, owner = Owner, mfa = {Module, Function, length(Args)}}.

%%% @doc
%%% Awaits a task reply and returns the result.
%%%
%%% This function waits for the task to complete and returns its result.
%%% If the task fails, this function will raise an error with the same
%%% reason as the task.
%%%
%%% A timeout, in milliseconds or :infinity, can be given with a default
%%% value of 5000. If the timeout is exceeded, await/2 will exit.
%%%
%%% This function can only be called by the process that started the task.
%%% @end
-spec await(task()) -> term().
await(Task) ->
    await(Task, 5000).

-spec await(task(), timeout()) -> term().
await(#task{ref = Ref, owner = Owner, pid = Pid}, Timeout) ->
    case self() of
        Owner ->
            receive
                {Ref, Result} ->
                    case Result of
                        {Class, Reason, Stacktrace} when Class =:= error; Class =:= exit; Class =:= throw ->
                            erlang:raise(Class, Reason, Stacktrace);
                        _ ->
                            Result
                    end;
                {'DOWN', _, process, Pid, Reason} when Pid =/= undefined ->
                    exit({Reason, {task, await, [#task{ref = Ref, owner = Owner, pid = Pid, mfa = {task, await, 2}}, Timeout]}})
            after Timeout ->
                exit({timeout, {task, await, [#task{ref = Ref, owner = Owner, pid = Pid, mfa = {task, await, 2}}, Timeout]}})
            end;
        _ ->
            error({not_owner, {task, await, [#task{ref = Ref, owner = Owner, pid = Pid, mfa = {task, await, 2}}, Timeout]}})
    end.

%%% @doc
%%% Yields to a task for a given time interval.
%%%
%%% This function checks if a task has completed within the given timeout.
%%% If the task completes, it returns {ok, Result}. If the timeout expires,
%%% it returns nil. If the task fails, it returns {exit, Reason}.
%%%
%%% Unlike await/2, this function does not raise an error if the task fails.
%%% Instead, it returns the error reason wrapped in an exit tuple.
%%%
%%% This function is useful when you want to periodically check on a task
%%% or when you want to handle task failures gracefully.
%%% @end
-spec yield(task()) -> {ok, term()} | {exit, term()} | nil.
yield(Task) ->
    yield(Task, 5000).

-spec yield(task(), timeout()) -> {ok, term()} | {exit, term()} | nil.
yield(#task{ref = Ref, owner = Owner, pid = Pid}, Timeout) ->
    case self() of
        Owner ->
            receive
                {Ref, Result} ->
                    case Result of
                        {Class, Reason, _Stacktrace} when Class =:= error; Class =:= exit; Class =:= throw ->
                            {exit, Reason};
                        _ ->
                            {ok, Result}
                    end;
                {'DOWN', _, process, Pid, Reason} when Pid =/= undefined ->
                    {exit, Reason}
            after Timeout ->
                nil
            end;
        _ ->
            error({not_owner, {task, yield, [#task{ref = Ref, owner = Owner, pid = Pid, mfa = {task, yield, 2}}, Timeout]}})
    end.

%%% @doc
%%% Starts a task as part of a supervision tree with the given function.
%%%
%%% This function spawns a new process as part of a supervision tree.
%%% The function should be a zero-arity anonymous function.
%%%
%%% This is used to start a statically supervised task under a supervision tree.
%%% Unlike async/1, this function does not return a task record and the result
%%% cannot be awaited. It's meant for fire-and-forget operations.
%%% @end
-spec start_link(task_fun()) -> {ok, pid()}.
start_link(Fun) when is_function(Fun, 0) ->
    start_link(erlang, apply, [Fun, []]).

%%% @doc
%%% Starts a task as part of a supervision tree with the given
%%% module, function, and arguments.
%%%
%%% This is used to start a statically supervised task under a supervision tree.
%%% Unlike async/3, this function does not return a task record and the result
%%% cannot be awaited. It's meant for fire-and-forget operations.
%%% @end
-spec start_link(module(), atom(), [term()]) -> {ok, pid()}.
start_link(Module, Function, Args) when is_atom(Module), is_atom(Function), is_list(Args) ->
    Owner = self(),
    Pid = spawn_link(fun() ->
        put('$initial_call', {Module, Function, length(Args)}),
        put('$ancestors', [Owner]),
        apply(Module, Function, Args)
    end),
    {ok, Pid}.

%%% @doc
%%% Starts a task with the given function.
%%%
%%% This function spawns a new process that is not linked to the caller.
%%% The function should be a zero-arity anonymous function.
%%%
%%% This should only be used when the task is used for side-effects
%%% (like I/O) and you have no interest in its results nor if it
%%% completes successfully.
%%% @end
-spec start(task_fun()) -> {ok, pid()}.
start(Fun) when is_function(Fun, 0) ->
    start(erlang, apply, [Fun, []]).

%%% @doc
%%% Starts a task with the given module, function, and arguments.
%%%
%%% This function spawns a new process that is not linked to the caller.
%%%
%%% This should only be used when the task is used for side-effects
%%% (like I/O) and you have no interest in its results nor if it
%%% completes successfully.
%%% @end
-spec start(module(), atom(), [term()]) -> {ok, pid()}.
start(Module, Function, Args) when is_atom(Module), is_atom(Function), is_list(Args) ->
    Pid = spawn(fun() ->
        put('$initial_call', {Module, Function, length(Args)}),
        apply(Module, Function, Args)
    end),
    {ok, Pid}.

%%% @doc
%%% Shuts down the task with the given reason.
%%%
%%% If the task is not linked, it will be killed immediately.
%%% If the task is linked and trapping exits, it will be asked to
%%% terminate with the given reason. If the task does not terminate
%%% within the timeout, it will be killed.
%%%
%%% Returns {ok, Result} if the task was already completed,
%%% {exit, Reason} if the task exited, or nil if the task
%%% was killed.
%%% @end
-spec shutdown(task()) -> {ok, term()} | {exit, term()} | nil.
shutdown(Task) ->
    shutdown(Task, 5000).

-spec shutdown(task(), timeout()) -> {ok, term()} | {exit, term()} | nil.
shutdown(#task{pid = undefined} = Task, Timeout) ->
    % Task was already completed, try to get the result
    yield(Task, Timeout);
shutdown(#task{pid = Pid, ref = Ref, owner = Owner} = Task, Timeout) ->
    case self() of
        Owner ->
            MonitorRef = monitor(process, Pid),
            exit(Pid, shutdown),
            receive
                {Ref, Result} ->
                    demonitor(MonitorRef, [flush]),
                    case Result of
                        {Class, Reason, _Stacktrace} when Class =:= error; Class =:= exit; Class =:= throw ->
                            {exit, Reason};
                        _ ->
                            {ok, Result}
                    end;
                {'DOWN', MonitorRef, process, Pid, Reason} ->
                    % Flush any pending task result message
                    receive
                        {Ref, _} -> ok
                    after 0 -> ok
                    end,
                    {exit, Reason}
            after Timeout ->
                demonitor(MonitorRef, [flush]),
                exit(Pid, kill),
                % Wait a bit for the kill to take effect
                receive
                    {'DOWN', _, process, Pid, _} -> ok
                after 1000 -> ok
                end,
                nil
            end;
        _ ->
            error({not_owner, {task, shutdown, [Task, Timeout]}})
    end.

%%% @doc
%%% Returns a task that is already completed with the given result.
%%%
%%% This function creates a task that immediately has a result available.
%%% It can be awaited or yielded like any other task.
%%%
%%% This is useful when you want to create a task that represents
%%% a computation that has already been completed, often for consistency
%%% in APIs that work with both synchronous and asynchronous operations.
%%% @end
-spec completed(term()) -> task().
completed(Result) ->
    Ref = make_ref(),
    Owner = self(),

    % Send the result immediately
    Owner ! {Ref, Result},

    #task{pid = undefined, ref = Ref, owner = Owner, mfa = {task, completed, 1}}.