-module(mongo@aggregation).
-compile([no_auto_import, nowarn_unused_vars]).

-export([aggregate/2, exec/1, stages/2, match/2, lookup/5, project/2, add_fields/2, sort/2, group/2, skip/2, limit/2]).
-export_type([pipeline/0, aggregate_option/0]).

-opaque pipeline() :: {pipeline,
        mongo@client:collection(),
        list(aggregate_option()),
        gleam@queue:queue(bson@value:value())}.

-type aggregate_option() :: {batch_size, integer()} |
    {'let', bson@value:value()}.

-spec aggregate(mongo@client:collection(), list(aggregate_option())) -> pipeline().
aggregate(Collection, Options) ->
    {pipeline, Collection, Options, gleam@queue:new()}.

-spec exec(pipeline()) -> {ok, mongo@cursor:cursor()} |
    {error, mongo@utils:mongo_error()}.
exec(Pipeline) ->
    Body = gleam@list:fold(
        erlang:element(3, Pipeline),
        [{<<"aggregate"/utf8>>,
                {str, erlang:element(3, erlang:element(2, Pipeline))}},
            {<<"pipeline"/utf8>>,
                begin
                    _pipe = erlang:element(4, Pipeline),
                    _pipe@1 = gleam@queue:to_list(_pipe),
                    {array, _pipe@1}
                end},
            {<<"cursor"/utf8>>, {document, []}}],
        fun(Acc, Opt) -> case Opt of
                {batch_size, Size} ->
                    gleam@list:key_set(
                        Acc,
                        <<"cursor"/utf8>>,
                        {document, [{<<"batchSize"/utf8>>, {int32, Size}}]}
                    );

                {'let', {document, Let_doc}} ->
                    gleam@list:key_set(Acc, <<"let"/utf8>>, {document, Let_doc})
            end end
    ),
    case mongo@client:execute(erlang:element(2, Pipeline), {document, Body}) of
        {ok, Result} ->
            [{<<"cursor"/utf8>>, {document, Result@1}}, {<<"ok"/utf8>>, Ok}] = Result,
            [{<<"firstBatch"/utf8>>, {array, Batch}},
                {<<"id"/utf8>>, {int64, Id}},
                {<<"ns"/utf8>>, _}] = Result@1,
            case Ok of
                {double, 1.0} ->
                    _pipe@2 = mongo@cursor:new(
                        erlang:element(2, Pipeline),
                        Id,
                        Batch
                    ),
                    {ok, _pipe@2};

                _ ->
                    {error, {mongo_error, -16, <<""/utf8>>, null}}
            end;

        {error, {Code, Msg}} ->
            {error, {mongo_error, Code, Msg, null}}
    end.

-spec append_stage(pipeline(), bson@value:value()) -> pipeline().
append_stage(Pipeline, Stage) ->
    {pipeline,
        erlang:element(2, Pipeline),
        erlang:element(3, Pipeline),
        begin
            _pipe = erlang:element(4, Pipeline),
            gleam@queue:push_back(_pipe, Stage)
        end}.

-spec stages(pipeline(), list(bson@value:value())) -> pipeline().
stages(Pipeline, Docs) ->
    gleam@list:fold(
        Docs,
        Pipeline,
        fun(New_pipeline, Current) -> append_stage(New_pipeline, Current) end
    ).

-spec match(pipeline(), bson@value:value()) -> pipeline().
match(Pipeline, Doc) ->
    append_stage(Pipeline, {document, [{<<"$match"/utf8>>, Doc}]}).

-spec lookup(pipeline(), binary(), binary(), binary(), binary()) -> pipeline().
lookup(Pipeline, From, Local_field, Foreign_field, Alias) ->
    append_stage(
        Pipeline,
        {document,
            [{<<"$lookup"/utf8>>,
                    {document,
                        [{<<"from"/utf8>>, {str, From}},
                            {<<"localField"/utf8>>, {str, Local_field}},
                            {<<"foreignField"/utf8>>, {str, Foreign_field}},
                            {<<"as"/utf8>>, {str, Alias}}]}}]}
    ).

-spec project(pipeline(), bson@value:value()) -> pipeline().
project(Pipeline, Doc) ->
    append_stage(Pipeline, {document, [{<<"$project"/utf8>>, Doc}]}).

-spec add_fields(pipeline(), bson@value:value()) -> pipeline().
add_fields(Pipeline, Doc) ->
    append_stage(Pipeline, {document, [{<<"$addFields"/utf8>>, Doc}]}).

-spec sort(pipeline(), bson@value:value()) -> pipeline().
sort(Pipeline, Doc) ->
    append_stage(Pipeline, {document, [{<<"$sort"/utf8>>, Doc}]}).

-spec group(pipeline(), bson@value:value()) -> pipeline().
group(Pipeline, Doc) ->
    append_stage(Pipeline, {document, [{<<"$group"/utf8>>, Doc}]}).

-spec skip(pipeline(), integer()) -> pipeline().
skip(Pipeline, Count) ->
    append_stage(Pipeline, {document, [{<<"$skip"/utf8>>, {int32, Count}}]}).

-spec limit(pipeline(), integer()) -> pipeline().
limit(Pipeline, Count) ->
    append_stage(Pipeline, {document, [{<<"$limit"/utf8>>, {int32, Count}}]}).