-module(gmser_dyn).

-export([ encode/1
        , encode/2
        , encode_typed/2
        , encode_typed/3
        , decode/1
        , decode/2 ]).

-export([ serialize/1
        , serialize/2
        , serialize_typed/2
        , serialize_typed/3
        , deserialize/1
        , deserialize/2 ]).

-export([ register_types/1
        , registered_types/0
        , revert_to_default_types/0
        , dynamic_types/0 ]).

-import(gmserialization, [ decode_field/2 ]).

-define(VSN, 1).

-include_lib("kernel/include/logger.hrl").

-ifdef(TEST).
-compile([export_all, nowarn_export_all]).
-include_lib("eunit/include/eunit.hrl").
-endif.

serialize(Term)                    -> rlp_encode(encode(Term)).
serialize(Term, Types)             -> rlp_encode(encode(Term, Types)).
serialize_typed(Type, Term)        -> rlp_encode(encode_typed(Type, Term)).
serialize_typed(Type, Term, Types) -> rlp_encode(encode_typed(Type, Term, Types)).

deserialize(Binary) -> decode(rlp_decode(Binary)).
deserialize(Binary, Types) -> decode(rlp_decode(Binary), Types).


encode(Term) ->
    encode(Term, registered_types()).

encode(Term, Types) ->
    encode(Term, vsn(Types), Types).

encode(Term, Vsn, Types) ->
    [ encode_basic(int, 0)
    , encode_basic(int, Vsn)
    , encode_(Term, Vsn, Types) ].

encode_typed(Type, Term) ->
    encode_typed(Type, Term, registered_types()).

encode_typed(Type, Term, Types) ->
    encode_typed(Type, Term, vsn(Types), Types).

encode_typed(Type, Term, Vsn, Types) ->
    [ encode_basic(int, 0)
    , encode_basic(int, Vsn)
    , encode_typed_(Type, Term, Vsn, Types) ].

decode(Fields) ->
    decode(Fields, registered_types()).

decode(Fields0, Types) ->
    case decode_tag_and_vsn(Fields0) of
        {0, Vsn, Fields} ->
            [Val] = decode_(Fields, Vsn, Types, []),
            Val;
        Other ->
            error({illegal_serialization, Other})
    end.

decode_tag_and_vsn([TagBin, VsnBin | Fields]) ->
    {decode_basic(int, TagBin),
     decode_basic(int, VsnBin),
     Fields}.

dynamic_types() ->
    #{ vsn   => ?VSN
     , codes =>
           #{ 248 => int
            , 249 => binary
            , 250 => bool
            , 251 => list
            , 252 => map
            , 253 => tuple
            , 254 => id
            , 255 => label }
     , rev =>
           #{ int    => 248
            , binary => 249
            , bool   => 250
            , list   => 251
            , map    => 252
            , tuple  => 253
            , id     => 254
            , label  => 255}
     , templates =>
           #{ int    => int
            , binary => binary
            , bool   => bool
            , list   => list
            , map    => map
            , tuple  => tuple
            , id     => id
            , label  => label
            }
       }.

vsn(Types) ->
    maps:get(vsn, Types, ?VSN).

register_types(Types) when is_map(Types) ->
    Codes = maps:get(codes, Types, #{}),
    Rev   = rev_codes(Codes),
    Templates = maps:get(templates, Types, #{}),
    #{codes := Codes0, rev := Rev0, templates := Templates0} =
        dynamic_types(),
    Merged = #{ codes => maps:merge(Codes0, Codes)
              , rev   => maps:merge(Rev0, Rev)
              , templates => maps:merge(Templates0, Templates) },
    assert_sizes(Merged),
    assert_mappings(Merged),
    persistent_term:put({?MODULE, types}, Merged).

revert_to_default_types() ->
    persistent_term:put({?MODULE, types}, dynamic_types()).

assert_sizes(#{codes := Codes, rev := Rev, templates := Ts} = Types) ->
    assert_sizes(map_size(Codes), map_size(Rev), map_size(Ts), Types).

assert_sizes(Sz, Sz, Sz, _) ->
    ok;
assert_sizes(Sz, RSz, Sz, Types) when RSz =/= Sz ->
    %% Wrong size reverse mapping must mean duplicate mappings
    %% We auto-generate the reverse-mappings, so we know there aren't
    %% too many of them
    ?LOG_ERROR("Reverse mapping size doesn't match codes size", []),
    Codes = maps:get(codes, Types),
    CodeVals = maps:values(Codes),
    Duplicates = CodeVals -- lists:usort(CodeVals),
    error({duplicate_mappings, Duplicates, Types});
assert_sizes(Sz, _, TSz, Types) when Sz > TSz ->
    ?LOG_ERROR("More codes than templates", []),
    Tags = maps:keys(maps:get(rev, Types)),
    Templates = maps:get(templates, Types),
    Missing = [T || T <- Tags,
                    not is_map_key(T, Templates)],
    error({missing_mappings, Missing, Types});
assert_sizes(Sz, _, TSz, Types) when TSz > Sz ->
    %% More mappings than codes. May not be horrible.
    %% We check that all codes have mappings elsewhere.
    ?LOG_WARNING("More templates than codes in ~p", [Types]),
    ok.

assert_mappings(#{rev := Rev, templates := Ts} = Types) ->
    Tags = maps:keys(Rev),
    case [T || T <- Tags,
               not is_map_key(T, Ts)] of
        [] ->
            ok;
        Missing ->
            ?LOG_ERROR("Missing templates for ~p", [Missing]),
            error({missing_templates, Missing, Types})
    end.

rev_codes(Codes) ->
    L = maps:to_list(Codes),
    maps:from_list([{V, K} || {K, V} <- L]).

registered_types() ->
    persistent_term:get({?MODULE, types}, dynamic_types()).

template(TagOrCode, Vsn, Types) ->
    {Tag, Template} = get_template(TagOrCode, Types),
    {Tag, dyn_template_(Template, Vsn)}.

get_template(Code, #{codes := Codes, templates := Ts}) when is_integer(Code) ->
    Tag = maps:get(Code, Codes),
    {Tag, maps:get(Tag, Ts)};
get_template(Tag, #{templates := Ts}) when is_atom(Tag) ->
    {Tag, maps:get(Tag, Ts)}.

dyn_template_(F, Vsn) ->
  if is_function(F, 0) -> F();
     is_function(F, 1) -> F(Vsn);
     true -> F
  end.

decode_(Fields, Vsn, Types, Acc) ->
    {_Tag, Term, Rest} = decode_field_(Fields, Vsn, Types),
    Acc1 = [Term | Acc],
    case Rest of
        [] ->
            lists:reverse(Acc1);
        _ ->
            decode_(Rest, Vsn, Types, Acc1)
    end.

decode_field_([H|T], Vsn, Types) ->
    {CodeBin, Field, Rest} =
        case H of
            [C, F] -> {C, F, T};
            C when is_binary(C) -> {C, hd(T), tl(T)}
        end,
    Code = decode_basic(int, CodeBin),
    {Tag, Template} = template(Code, Vsn, Types),
    %% [Fld|Rest] = Fields,
    Val = decode_from_template(Template, Field, Vsn, Types),
    {Tag, Val, Rest}.

encode_(Term, Vsn, Types) ->
    encode_(Term, true, Vsn, Types).

encode_(Term, Emit, Vsn, Types) ->
    {Tag, Template} = auto_template(Term),
    Enc = encode_from_template(Template, Term, Vsn, Types),
    if Emit ->
            [emit_code(Tag, Types), Enc];
       true ->
            Enc
    end.

encode_typed_(Code, Term, Vsn, #{codes := Codes} = Types) when is_map_key(Code, Codes) ->
    {_Tag, Template} = template(Code, Vsn, Types),
    [encode_basic(int, Code), encode_from_template(Template, Term, Vsn, Types)];
encode_typed_(Tag, Term, Vsn, #{templates := Ts} = Types) when is_map_key(Tag, Ts) ->
    Template = maps:get(Tag, Ts),
    [emit_code(Tag, Types), encode_from_template(Template, Term, Vsn, Types)];
encode_typed_(MaybeTemplate, Term, Vsn, Types) ->
    encode_maybe_template(MaybeTemplate, Term, Vsn, Types).

encode_maybe_template(Pat, Term, Vsn, Types) when is_list(Pat);
                                                  is_tuple(Pat);
                                                  is_map(Pat) ->
    {Tag, _} = auto_template(Pat),
    [emit_code(Tag, Types),
     encode_from_template(Pat, Term, Vsn, Types)];
encode_maybe_template(Other, Term, _Vsn, _Types) ->
    error({illegal_template, Other, Term}).

auto_template({id,Tag,V}) when Tag == account
                               ; Tag == name
                               ; Tag == commitment
                               ; Tag == contract
                               ; Tag == channel
                               ; Tag == associate_chain
                               ; Tag == entry ->
    if is_binary(V) -> {id, id};
       true ->
            %% close, but no cigar
            {tuple, tuple}
    end;
auto_template(T) ->
    if is_map(T)     -> {map, map};
       is_list(T)    -> {list, list};
       is_tuple(T)   -> {tuple, tuple};
       is_binary(T)  -> {binary, binary};
       is_boolean(T) -> {bool, bool};
       is_atom(T)    -> {label, label};   % binary_to_existing_atom()
       is_integer(T),
       T >= 0        -> {int, int};
       true ->
            error(invalid_type)
    end.

decode_from_template(list, Fld, Vsn, Types) ->
    decode_(Fld, Vsn, Types, []);
decode_from_template(map, Fld, Vsn, Types) ->
    TupleFields = [F || F <- Fld],
    Items = [decode_from_template(tuple, T, Vsn, Types)
             || T <- TupleFields],
    maps:from_list(Items);
decode_from_template(tuple, Fld, Vsn, Types) ->
    Items = decode_(Fld, Vsn, Types, []),
    list_to_tuple(Items);
decode_from_template([Type], Fields, Vsn, Types) ->
    [decode_from_template(Type, F, Vsn, Types)
     || F <- Fields];
decode_from_template(Type, V, Vsn, Types) when is_list(Type), is_list(V) ->
    decode_fields(Type, V, Vsn, Types);
decode_from_template(Type, V, Vsn, Types) when is_tuple(Type), is_list(V) ->
    Zipped = lists:zip(tuple_to_list(Type), V),
    Items = [decode_from_template(T1, V1, Vsn, Types) || {T1, V1} <- Zipped],
    list_to_tuple(Items);
decode_from_template(Type, Fld, _, _) when Type == int
                                           ; Type == binary
                                           ; Type == bool
                                           ; Type == id
                                           ; Type == label ->
    decode_basic(Type, Fld).    

encode_from_template(Type, V, Vsn, Types) ->
    encode_from_template(Type, V, true, Vsn, Types).

encode_from_template(list, L, _, Vsn, Types) when is_list(L) ->
    [encode_(V, Vsn, Types) || V <- L];
encode_from_template(map, M, _, Vsn, Types) when is_map(M) ->
    [encode_({K,V}, false, Vsn, Types)
     || {K, V} <- lists:sort(maps:to_list(M))];
encode_from_template(tuple, T, _, Vsn, Types) when is_tuple(T) ->
    [encode_(V, Vsn, Types) || V <- tuple_to_list(T)];
encode_from_template(T, V, _, Vsn, Types) when tuple_size(T) =:= tuple_size(V) ->
    Zipped = lists:zip(tuple_to_list(T), tuple_to_list(V)),
    [encode_from_template(T1, V1, false, Vsn, Types) || {T1, V1} <- Zipped];
encode_from_template([Type], List, _, Vsn, Types) ->
    [encode_from_template(Type, V, false, Vsn, Types) || V <- List];
encode_from_template(Type, List, _, Vsn, Types) when is_list(Type), is_list(List) ->
    encode_fields(Type, List, Vsn, Types);
encode_from_template(Type, V, _, _, _Types) when Type == id
                                                 ; Type == binary
                                                 ; Type == bool
                                                 ; Type == int
                                                 ; Type == label ->
    encode_basic(Type, V);
encode_from_template(Type, V, _, _, _) ->
    error({illegal, Type, V}).


%% Basically, dynamically encoding a statically defined object
encode_fields([{Field, Type}|TypesLeft],
              [{Field,  Val}|FieldsLeft], Vsn, Types) ->
    [ encode_from_template(Type, Val, Vsn, Types)
    | encode_fields(TypesLeft, FieldsLeft, Vsn, Types)];
encode_fields([{_Field, Type}|TypesLeft],
              [Val           |FieldsLeft], Vsn, Types) ->
    %% Not sure if we want to try this ...
    [ encode_from_template(Type, Val, Vsn, Types)
    | encode_fields(TypesLeft, FieldsLeft, Vsn, Types)];
encode_fields([Type|TypesLeft],
              [Val |FieldsLeft], Vsn, Types) when is_atom(Type) ->
    %% Not sure about this either ...
    [ encode_from_template(Type, Val, Vsn, Types)
    | encode_fields(TypesLeft, FieldsLeft, Vsn, Types)];
encode_fields([], [], _, _) ->
    [].

decode_fields([{Tag, Type}|TypesLeft],
              [Field      |FieldsLeft], Vsn, Types) ->
    
    [ {Tag, decode_from_template(Type, Field, Vsn, Types)}
    | decode_fields(TypesLeft, FieldsLeft, Vsn, Types)];
decode_fields([], [], _, _) ->
    [].

emit_code(Tag, #{rev := Tags}) ->
    encode_basic(int, maps:get(Tag, Tags)).

decode_basic(label, Fld) ->
    binary_to_existing_atom(decode_basic(binary, Fld), utf8);
decode_basic(Type, Fld) ->
    gmserialization:decode_field(Type, Fld).

encode_basic(label, A) when is_atom(A) ->
    encode_basic(binary, atom_to_binary(A, utf8));
encode_basic(Type, Fld) ->
    gmserialization:encode_field(Type, Fld).

rlp_decode(Bin) ->
    gmser_rlp:decode(Bin).

rlp_encode(Fields) ->
    gmser_rlp:encode(Fields).


-ifdef(TEST).

trace() ->
    dbg:tracer(),
    dbg:tpl(?MODULE, x),
    dbg:p(all, [c]).

notrace() ->
    dbg:ctpl('_'),
    dbg:stop().

round_trip_test_() ->
    [?_test(t_round_trip(T)) ||
        T <- t_sample_types()
    ].

t_sample_types() ->
    [ 5
    , <<"a">>
    , [1,2,3]
    , {<<"a">>,1}
    , #{<<"a">> => 1}
    , [#{1 => <<"c">>, [17] => true}]
    , true
    ].

user_types_test_() ->
    {foreach,
     fun() ->
             revert_to_default_types()
     end,
     fun(_) ->
             revert_to_default_types()
     end,
     [ ?_test(t_reg_typed_tuple())
     , ?_test(t_reg_chain_objects_array())
     ]}.

t_round_trip(T) ->
    ?debugVal(T),
    ?assertMatch({T, T}, {T, decode(encode(T))}).

t_reg_typed_tuple() ->
    Type = {int, int, int},
    MyTypes = #{ codes => #{ 1001 => int_tup3 }
               , templates => #{ int_tup3 => Type }
               },
    register_types(MyTypes),
    GoodTerm = {2,3,4},
    ?debugFmt("Type: ~p, GoodTerm = ~p", [Type, GoodTerm]),
    Enc = encode_typed(int_tup3, GoodTerm),
    GoodTerm = decode(Enc),
    t_bad_typed_encode(int_tup3, {1,2,<<"a">>}, {illegal,int,<<"a">>}),
    t_bad_typed_encode(int_tup3, {1,2,3,4}, {illegal, {int,int,int}, {1,2,3,4}}).

t_bad_typed_encode(Type, Term, Error) ->
    try encode_typed(Type, Term),
         error({expected_error, Error})
    catch
        error:Error ->
            ok
    end.

t_reg_chain_objects_array() ->
    Template = [{foo, {int, binary}},  {bar, [{int, int}]}, {baz, {int}}],
    ?debugFmt("Template = ~p", [Template]),
    MyTypes = #{ codes => #{ 1002 => coa }
               , templates => #{ coa => Template } },
    register_types(MyTypes),
    Values   = [{foo, {1, <<"foo">>}}, {bar, [{1, 2}, {3, 4}, {5, 6}]}, {baz, {1}}],
    ?debugFmt("Values = ~p", [Values]),
    Enc = encode_typed(coa, Values),
    Values = decode(Enc).

-endif.