%%% Copyright (c) 2014, NORDUnet A/S.
%%% See LICENSE for licensing information.
%%%

%% tree storage, TODO: docdoc

-module(ts).
-include_lib("eunit/include/eunit.hrl").
-export_type([tree_store/0]).
-export([new/0, add/3, delete/2, retrieve/2, count/2]).

%% FIXME: Keep the entries in binaries instead of lists? Hashes do
%% have fixed lenght.
-record(tree_store, {layers :: list()}). % orddict of lists, keyed on layer.
-type tree_store() :: #tree_store{}.

%%%%%%%%%%%%%%%%%%%%
%% Public.
new() ->
    #tree_store{layers = orddict:new()}.

-spec add(tree_store(), non_neg_integer(), binary()) -> tree_store().
add(S = #tree_store{layers = Layers}, Layer, Entry) ->
    {NewLayers, List} = layer(Layers, rw, Layer),
    NewList = [Entry | List],
    S#tree_store{layers = orddict:store(Layer, NewList, NewLayers)}.

-spec delete(tree_store(), non_neg_integer()) -> tree_store().
delete(S = #tree_store{layers = Layers}, Layer) ->
    List = layer(Layers, ro, Layer),
    [_ | NewList] = List,
    S#tree_store{layers = orddict:store(Layer, NewList, Layers)}.

-spec retrieve(tree_store(), tuple()) -> binary() | undefined.
retrieve(#tree_store{layers = Layers}, {Layer, Index}) ->
    List = layer(Layers, ro, Layer),
    Len = length(List),
    case Index < Len of
        true -> lists:nth(Len - Index, List);
        false -> undefined
    end.

-spec count(tree_store(), non_neg_integer()) -> non_neg_integer().
count(#tree_store{layers = Layers}, Layer) ->
    length(layer(Layers, ro, Layer)).

%%%%%%%%%%%%%%%%%%%%
%% Private.
-spec layer(list(), rw | ro, non_neg_integer()) -> list() | {list(), list()}.
layer(Layers, Access, Layer) ->
    case Access of
        rw ->
            case orddict:find(Layer, Layers) of
                error -> {orddict:store(Layer, [], Layers), []};
                {ok, List} -> {Layers, List}
            end;
        ro ->
            case orddict:find(Layer, Layers) of
                error -> [];
                {ok, List} -> List
            end
    end.

%%%%%%%%%%%%%%%%%%%%
%% Testing ts.
add_retrieve_test_() ->
    S = new(),
    S0 = add(S, 0, <<00>>),
    S1 = add(S0, 0, <<01>>),
    S2 = add(S1, 1, <<10>>),
    [?_assertEqual(2, count(S2, 0)),
     ?_assertEqual(1, count(S2, 1)),
     ?_assertEqual(<<10>>, retrieve(S2, {1,0})),
     ?_assertEqual(<<00>>, retrieve(S2, {0,0})),
     ?_assertEqual(<<01>>, retrieve(S2, {0,1})),
     ?_assertEqual(<<00>>, retrieve(S2, {0,0})),
     ?_assertEqual(undefined, retrieve(S2, {3,0})),
     ?_assertEqual(undefined, retrieve(S1, {1,0}))].

delete_test_() ->
    S = new(),
    S0 = add(S, 0, <<00>>),
    S1 = add(S0, 0, <<01>>),
    S2 = add(S1, 1, <<10>>),
    S3 = delete(S2, 0),
    S4 = add(S3, 0, <<99>>),
    [?_assertEqual(2, count(S2, 0)),
     ?_assertEqual(1, count(S3, 0)),
     ?_assertEqual(2, count(S4, 0)),
     ?_assertEqual(retrieve(S2, {0,1}), <<01>>),
     ?_assertEqual(retrieve(S3, {0,1}), undefined),
     ?_assertEqual(retrieve(S3, {0,0}), <<00>>),
     ?_assertEqual(retrieve(S4, {0,1}), <<99>>)].