%%%-------------------------------------------------------------------
%%% @copyright (C) 2016, AdRoll
%%% @doc
%%%
%%%     Kinesis record aggregator.
%%%
%%%     Follows the KPL aggregated record format:
%%%     https://github.com/awslabs/amazon-kinesis-producer/blob/master/aggregation-format.md
%%%
%%%     This is an Erlang port of the aggregation functionality from:
%%%     https://pypi.python.org/pypi/aws_kinesis_agg/1.0.0
%%%
%%%     Creating a new aggregator:
%%%
%%%         Agg = kpl_agg:new()
%%%
%%%     Adding user records to an aggregator (the aggregator will emit an
%%%     aggregated record when it is full):
%%%
%%%         case kpl_agg:add(Agg, Record) of
%%%             {undefined, NewAgg} -> ...
%%%             {FullAggRecord, NewAgg} -> ...
%%%         end
%%%
%%%     You can also use kpl:add_all to add multiple records at once. A
%%%     <pre>Record</pre> is a {PartitionKey, Data} tuple or a {PartitionKey, Data,
%%%     ExplicitHashKey} tuple.
%%%
%%%     Getting the current aggregated record (e.g. to get the last aggregated
%%%     record when you have no more user records to add):
%%%
%%%         case kpl_agg:finish(Agg) of
%%%             {undefined, Agg} -> ...
%%%             {AggRecord, NewAgg} -> ...
%%%         end
%%%
%%%     The result currently uses a non-standard magic prefix to prevent the KCL from
%%%     deaggregating the record automatically.  To use compression, instantiate the
%%%     aggregator using kpl_agg:new(true), which uses another
%%%     non-standard magic prefix.
%%%
%%% @end
%%% Created: 12 Dec 2016 by Constantin Berzan <constantin.berzan@adroll.com>
%%%-------------------------------------------------------------------
-module(kpl_agg).

%% API
-export([new/0, new/1, count/1, size_bytes/1, finish/1, add/2, add_all/2]).

-define(MD5_DIGEST_BYTES, 16).
%% From http://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecord.html:
-define(KINESIS_MAX_BYTES_PER_RECORD, 1 bsl 20).

-include("erlmld.hrl").
-include("kpl_agg_pb.hrl").

%% A set of keys, mapping each key to a unique index.
-record(keyset,
        {rev_keys = [] :: [binary()],    %% list of known keys in reverse order
         rev_keys_length = 0 :: non_neg_integer(), %% length of the rev_keys list
         key_to_index = maps:new() :: map()}). %% maps each known key to a 0-based index
%% Internal state of a record aggregator. It stores an aggregated record that
%% is "in progress", i.e. it is possible to add more user records to it.
-record(state,
        {num_user_records = 0 :: non_neg_integer(),
         agg_size_bytes = 0 :: non_neg_integer(),
         %% The aggregated record's PartitionKey and ExplicitHashKey are the
         %% PartitionKey and ExplicitHashKey of the first user record added.
         agg_partition_key = undefined :: undefined | binary(),
         agg_explicit_hash_key = undefined :: undefined | binary(),
         %% Keys seen in the user records added so far.
         partition_keyset = #keyset{} :: #keyset{},
         explicit_hash_keyset = #keyset{} :: #keyset{},
         %% List if user records added so far, in reverse order.
         rev_records = [] :: [#'Record'{}],
         should_deflate = false}).

%%%===================================================================
%%% API
%%%===================================================================

new() ->
    new(false).

new(ShouldDeflate) ->
    #state{should_deflate = ShouldDeflate}.

count(#state{num_user_records = Num} = _State) ->
    Num.

size_bytes(#state{agg_size_bytes = Size, agg_partition_key = PK} = _State) ->
    PKSize =
        case PK of
            undefined ->
                0;
            _ ->
                byte_size(PK)
        end,
    byte_size(?KPL_AGG_MAGIC)
    + Size
    + ?MD5_DIGEST_BYTES
    + PKSize
    + byte_size(kpl_agg_pb:encode_msg(#'AggregatedRecord'{})).

finish(#state{num_user_records = 0} = State) ->
    {undefined, State};
finish(#state{agg_partition_key = AggPK,
              agg_explicit_hash_key = AggEHK,
              should_deflate = ShouldDeflate} =
           State) ->
    AggRecord = {AggPK, serialize_data(State, ShouldDeflate), AggEHK},
    {AggRecord, new(ShouldDeflate)}.

add(State, {PartitionKey, Data} = _Record) ->
    add(State, {PartitionKey, Data, undefined});
add(State, {PartitionKey, Data, ExplicitHashKey} = _Record) ->
    case {calc_record_size(State, PartitionKey, Data, ExplicitHashKey), size_bytes(State)} of
        {RecSize, _} when RecSize > ?KINESIS_MAX_BYTES_PER_RECORD ->
            error("input record too large to fit in a single Kinesis record");
        {RecSize, CurSize} when RecSize + CurSize > ?KINESIS_MAX_BYTES_PER_RECORD ->
            {FullRecord, State1} = finish(State),
            State2 = add_record(State1, PartitionKey, Data, ExplicitHashKey, RecSize),
            {FullRecord, State2};
        {RecSize, _} ->
            State1 = add_record(State, PartitionKey, Data, ExplicitHashKey, RecSize),
            %% fixme; make size calculations more accurate
            case size_bytes(State1) > ?KINESIS_MAX_BYTES_PER_RECORD - 64 of
                true ->
                    %% size estimate is almost the limit, finish & retry:
                    {FullRecord, State2} = finish(State),
                    State3 = add_record(State2, PartitionKey, Data, ExplicitHashKey, RecSize),
                    {FullRecord, State3};
                false ->
                    {undefined, State1}
            end
    end.

add_all(State, Records) ->
    {RevAggRecords, NState} =
        lists:foldl(fun(Record, {RevAggRecords, Agg}) ->
                       case add(Agg, Record) of
                           {undefined, NewAgg} ->
                               {RevAggRecords, NewAgg};
                           {AggRecord, NewAgg} ->
                               {[AggRecord | RevAggRecords], NewAgg}
                       end
                    end,
                    {[], State},
                    Records),
    {lists:reverse(RevAggRecords), NState}.

%%%===================================================================
%%% Internal functions
%%%===================================================================

%% Calculate how many extra bytes the given user record would take, when added
%% to the current aggregated record. This calculation has to know about KPL and
%% Protobuf internals.
calc_record_size(#state{partition_keyset = PartitionKeySet,
                        explicit_hash_keyset = ExplicitHashKeySet} =
                     _State,
                 PartitionKey,
                 Data,
                 ExplicitHashKey) ->
    %% How much space we need for the PK:
    PKLength = byte_size(PartitionKey),
    PKSize =
        case is_key(PartitionKey, PartitionKeySet) of
            true ->
                0;
            false ->
                1 + varint_size(PKLength) + PKLength
        end,

    %% How much space we need for the EHK:
    EHKSize =
        case ExplicitHashKey of
            undefined ->
                0;
            _ ->
                EHKLength = byte_size(ExplicitHashKey),
                case is_key(ExplicitHashKey, ExplicitHashKeySet) of
                    true ->
                        0;
                    false ->
                        1 + varint_size(EHKLength) + EHKLength
                end
        end,

    %% How much space we need for the inner record:
    PKIndexSize = 1 + varint_size(potential_index(PartitionKey, PartitionKeySet)),
    EHKIndexSize =
        case ExplicitHashKey of
            undefined ->
                0;
            _ ->
                1 + varint_size(potential_index(ExplicitHashKey, ExplicitHashKeySet))
        end,
    DataLength = byte_size(Data),
    DataSize = 1 + varint_size(DataLength) + DataLength,
    InnerSize = PKIndexSize + EHKIndexSize + DataSize,

    %% How much space we need for the entire record:
    PKSize + EHKSize + 1 + varint_size(InnerSize) + InnerSize.

%% Calculate how many bytes are needed to represent the given integer in a
%% Protobuf message.
varint_size(Integer) when Integer >= 0 ->
    NumBits = max(num_bits(Integer, 0), 1),
    (NumBits + 6) div 7.

%% Recursively compute the number of bits needed to represent an integer.
num_bits(0, Acc) ->
    Acc;
num_bits(Integer, Acc) when Integer >= 0 ->
    num_bits(Integer bsr 1, Acc + 1).

%% Helper for add; do not use directly.
add_record(#state{partition_keyset = PKSet,
                  explicit_hash_keyset = EHKSet,
                  rev_records = RevRecords,
                  num_user_records = NumUserRecords,
                  agg_size_bytes = AggSize,
                  agg_partition_key = AggPK,
                  agg_explicit_hash_key = AggEHK} =
               State,
           PartitionKey,
           Data,
           ExplicitHashKey,
           NewRecordSize) ->
    {PKIndex, NewPKSet} = get_or_add_key(PartitionKey, PKSet),
    {EHKIndex, NewEHKSet} = get_or_add_key(ExplicitHashKey, EHKSet),
    NewRecord =
        #'Record'{partition_key_index = PKIndex,
                  explicit_hash_key_index = EHKIndex,
                  data = Data},
    State#state{partition_keyset = NewPKSet,
                explicit_hash_keyset = NewEHKSet,
                rev_records = [NewRecord | RevRecords],
                num_user_records = 1 + NumUserRecords,
                agg_size_bytes = NewRecordSize + AggSize,
                agg_partition_key = first_defined(AggPK, PartitionKey),
                agg_explicit_hash_key = first_defined(AggEHK, ExplicitHashKey)}.

first_defined(undefined, Second) ->
    Second;
first_defined(First, _) ->
    First.

serialize_data(#state{partition_keyset = PKSet,
                      explicit_hash_keyset = EHKSet,
                      rev_records = RevRecords} =
                   _State,
               ShouldDeflate) ->
    ProtobufMessage =
        #'AggregatedRecord'{partition_key_table = key_list(PKSet),
                            explicit_hash_key_table = key_list(EHKSet),
                            records = lists:reverse(RevRecords)},
    SerializedData = kpl_agg_pb:encode_msg(ProtobufMessage),
    Checksum = crypto:hash(md5, SerializedData),
    case ShouldDeflate of
        true ->
            <<?KPL_AGG_MAGIC_DEFLATED/binary,
              (zlib:compress(<<SerializedData/binary, Checksum/binary>>))/binary>>;
        false ->
            <<?KPL_AGG_MAGIC/binary, SerializedData/binary, Checksum/binary>>
    end.

%%%===================================================================
%%% Internal functions for keysets
%%%===================================================================

is_key(Key, #keyset{key_to_index = KeyToIndex} = _KeySet) ->
    maps:is_key(Key, KeyToIndex).

get_or_add_key(undefined, KeySet) ->
    {undefined, KeySet};
get_or_add_key(Key,
               #keyset{rev_keys = RevKeys,
                       rev_keys_length = Length,
                       key_to_index = KeyToIndex} =
                   KeySet) ->
    case maps:get(Key, KeyToIndex, not_found) of
        not_found ->
            NewKeySet =
                KeySet#keyset{rev_keys = [Key | RevKeys],
                              rev_keys_length = Length + 1,
                              key_to_index = maps:put(Key, Length, KeyToIndex)},
            {Length, NewKeySet};
        Index ->
            {Index, KeySet}
    end.

potential_index(Key,
                #keyset{rev_keys_length = Length, key_to_index = KeyToIndex} = _KeySet) ->
    case maps:get(Key, KeyToIndex, not_found) of
        not_found ->
            Length;
        Index ->
            Index
    end.

key_list(#keyset{rev_keys = RevKeys} = _KeySet) ->
    lists:reverse(RevKeys).

%%%===================================================================
%%% TESTS
%%%===================================================================

-ifdef(TEST).

-include_lib("eunit/include/eunit.hrl").

varint_size_test() ->
    %% Reference values obtained using
    %% aws_kinesis_agg.aggregator._calculate_varint_size().
    ?assertEqual(1, varint_size(0)),
    ?assertEqual(1, varint_size(1)),
    ?assertEqual(1, varint_size(127)),
    ?assertEqual(2, varint_size(128)),
    ?assertEqual(4, varint_size(9999999)),
    ?assertEqual(6, varint_size(999999999999)),
    ok.

keyset_test() ->
    KeySet0 = #keyset{},
    ?assertEqual([], key_list(KeySet0)),
    ?assertEqual(false, is_key(<<"foo">>, KeySet0)),
    ?assertEqual(0, potential_index(<<"foo">>, KeySet0)),

    {0, KeySet1} = get_or_add_key(<<"foo">>, KeySet0),
    ?assertEqual([<<"foo">>], key_list(KeySet1)),
    ?assertEqual(true, is_key(<<"foo">>, KeySet1)),
    {0, KeySet1} = get_or_add_key(<<"foo">>, KeySet1),
    ?assertEqual(1, potential_index(<<"bar">>, KeySet1)),

    {1, KeySet2} = get_or_add_key(<<"bar">>, KeySet1),
    ?assertEqual([<<"foo">>, <<"bar">>], key_list(KeySet2)),
    ?assertEqual(true, is_key(<<"foo">>, KeySet2)),
    ?assertEqual(true, is_key(<<"bar">>, KeySet2)),
    {0, KeySet2} = get_or_add_key(<<"foo">>, KeySet2),
    {1, KeySet2} = get_or_add_key(<<"bar">>, KeySet2),
    ?assertEqual(2, potential_index(<<"boom">>, KeySet2)),

    ok.

empty_aggregator_test() ->
    Agg = new(),
    ?assertEqual(0, count(Agg)),
    ?assertEqual(4 + 16, size_bytes(Agg)),  % magic and md5
    {undefined, Agg} = finish(Agg),
    ok.

simple_aggregation_test() ->
    Agg0 = new(),
    {undefined, Agg1} = add(Agg0, {<<"pk1">>, <<"data1">>, <<"ehk1">>}),
    {undefined, Agg2} = add(Agg1, {<<"pk2">>, <<"data2">>, <<"ehk2">>}),
    {AggRecord, Agg3} = finish(Agg2),
    ?assertEqual(0, count(Agg3)),
    %% Reference values obtained using priv/kpl_agg_tests_helper.py.
    RefPK = <<"pk1">>,
    RefEHK = <<"ehk1">>,
    RefData =
        <<?KPL_AGG_MAGIC/binary, 10, 3, 112, 107, 49, 10, 3, 112, 107, 50, 18, 4, 101, 104, 107,
          49, 18, 4, 101, 104, 107, 50, 26, 11, 8, 0, 16, 0, 26, 5, 100, 97, 116, 97, 49, 26, 11, 8,
          1, 16, 1, 26, 5, 100, 97, 116, 97, 50, 244, 41, 93, 155, 173, 190, 58, 30, 240, 223, 216,
          8, 26, 205, 86, 4>>,
    ?assertEqual({RefPK, RefData, RefEHK}, AggRecord),
    ok.

aggregate_many(Records) ->
    {AggRecords, Agg} = add_all(new(), Records),
    case finish(Agg) of
        {undefined, _} ->
            AggRecords;
        {LastAggRecord, _} ->
            AggRecords ++ [LastAggRecord]
    end.

shared_keys_test() ->
    [AggRecord] =
        aggregate_many([{<<"alpha">>, <<"data1">>, <<"zulu">>},
                        {<<"beta">>, <<"data2">>, <<"yankee">>},
                        {<<"alpha">>, <<"data3">>, <<"xray">>},
                        {<<"charlie">>, <<"data4">>, <<"yankee">>},
                        {<<"beta">>, <<"data5">>, <<"zulu">>}]),
    %% Reference values obtained using priv/kpl_agg_tests_helper.py.
    RefPK = <<"alpha">>,
    RefEHK = <<"zulu">>,
    RefData =
        <<?KPL_AGG_MAGIC/binary, 10, 5, 97, 108, 112, 104, 97, 10, 4, 98, 101, 116, 97, 10, 7, 99,
          104, 97, 114, 108, 105, 101, 18, 4, 122, 117, 108, 117, 18, 6, 121, 97, 110, 107, 101,
          101, 18, 4, 120, 114, 97, 121, 26, 11, 8, 0, 16, 0, 26, 5, 100, 97, 116, 97, 49, 26, 11,
          8, 1, 16, 1, 26, 5, 100, 97, 116, 97, 50, 26, 11, 8, 0, 16, 2, 26, 5, 100, 97, 116, 97,
          51, 26, 11, 8, 2, 16, 1, 26, 5, 100, 97, 116, 97, 52, 26, 11, 8, 1, 16, 0, 26, 5, 100, 97,
          116, 97, 53, 78, 67, 160, 206, 22, 1, 33, 154, 3, 6, 110, 235, 9, 229, 53, 100>>,
    ?assertEqual({RefPK, RefData, RefEHK}, AggRecord),
    ok.

record_fullness_test() ->
    Data1 = list_to_binary(["X" || _ <- lists:seq(1, 500000)]),
    Data2 = list_to_binary(["Y" || _ <- lists:seq(1, 600000)]),
    Data3 = list_to_binary(["Z" || _ <- lists:seq(1, 200000)]),

    Agg0 = new(),
    {undefined, Agg1} = add(Agg0, {<<"pk1">>, Data1, <<"ehk1">>}),
    {{AggPK1, _AggData1, AggEHK1}, Agg2} = add(Agg1, {<<"pk2">>, Data2, <<"ehk2">>}),
    {undefined, Agg3} = add(Agg2, {<<"pk3">>, Data3, <<"ehk3">>}),
    {{AggPK2, _AggData2, AggEHK2}, _} = finish(Agg3),

    %% Reference values obtained using priv/kpl_agg_tests_helper.py.
    % fixme; these comparisons will fail as long as we're using the wrong kpl magic.
    %RefChecksum1 = <<198,6,88,216,8,244,159,59,223,14,247,208,138,137,64,118>>,
    %RefChecksum2 = <<89,148,130,126,150,23,148,18,38,230,176,182,93,186,150,69>>,
    ?assertEqual(<<"pk1">>, AggPK1),
    ?assertEqual(<<"ehk1">>, AggEHK1),
    %?assertEqual(RefChecksum1, crypto:hash(md5, AggData1)),
    ?assertEqual(<<"pk2">>, AggPK2),
    ?assertEqual(<<"ehk2">>, AggEHK2),
    %?assertEqual(RefChecksum2, crypto:hash(md5, AggData2)),
    ok.

full_record_test() ->
    Fill =
        fun F(Acc) ->
                PK = integer_to_binary(rand:uniform(1000)),
                Data =
                    << <<(integer_to_binary(rand:uniform(128)))/binary>>
                       || _ <- lists:seq(1, 1 + rand:uniform(1000)) >>,
                case add(Acc, {PK, Data}) of
                    {undefined, NAcc} ->
                        F(NAcc);
                    {Full, _} ->
                        Full
                end
        end,
    {PK, Data, _} = Fill(new()),
    Total = byte_size(PK) + byte_size(Data),
    ?assert(Total =< ?KINESIS_MAX_BYTES_PER_RECORD),
    ?assert(Total >= ?KINESIS_MAX_BYTES_PER_RECORD - 2048).

deflate_test() ->
    Agg0 = new(true),
    {undefined, Agg1} = add(Agg0, {<<"pk1">>, <<"data1">>, <<"ehk1">>}),
    {{_, Data, _}, _} = finish(Agg1),
    <<Magic:4/binary, Deflated/binary>> = Data,
    ?assertEqual(?KPL_AGG_MAGIC_DEFLATED, Magic),
    Inflated = zlib:uncompress(Deflated),
    ProtoMsg = binary:part(Inflated, 0, size(Inflated) - 16),
    Checksum = binary:part(Inflated, size(Inflated), -16),
    ?assertEqual(Checksum, crypto:hash(md5, ProtoMsg)),
    #'AggregatedRecord'{records = [R1]} = kpl_agg_pb:decode_msg(ProtoMsg, 'AggregatedRecord'),
    #'Record'{data = RecordData} = R1,
    ?assertEqual(<<"data1">>, RecordData).

-endif.