% Licensed under the Apache License, Version 2.0 (the "License"); you may not
% use this file except in compliance with the License.  You may obtain a copy of
% the License at
%
%   http://www.apache.org/licenses/LICENSE-2.0
%
% Unless required by applicable law or agreed to in writing, software
% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
% License for the specific language governing permissions and limitations under
% the License.

-module(couch_key_tree).

-export([merge/2, find_missing/2, get_key_leafs/2, get_full_key_paths/2, get/2]).
-export([map/2, get_all_leafs/1, get_leaf_keys/1, count_leafs/1, remove_leafs/2,get_all_leafs_full/1]).

% a key tree looks like this:
% Tree -> [] or [{Key, Value, ChildTree} | SiblingTree]
% ChildTree -> Tree
% SiblingTree -> [] or [{SiblingKey, Value, Tree} | Tree]
% And each Key < SiblingKey



% key tree functions

% When the same key is found in the trees, the value in tree B is discarded.
merge([], B) ->
    B;
merge(A, []) ->
    A;
merge([ATree | ANextTree], [BTree | BNextTree]) ->
    {AKey, AValue, ASubTree} = ATree,
    {BKey, _BValue, BSubTree} = BTree,
    if
    AKey == BKey ->
        %same key
        MergedSubTree = merge(ASubTree, BSubTree),
        MergedNextTree = merge(ANextTree, BNextTree),
        [{AKey, AValue, MergedSubTree} | MergedNextTree];
    AKey < BKey ->
        [ATree | merge(ANextTree, [BTree | BNextTree])];
    true ->
        [BTree | merge([ATree | ANextTree], BNextTree)]
    end.

find_missing(_Tree, []) ->
    [];
find_missing([], Keys) ->
    Keys;
find_missing([{Key, _, SubTree} | RestTree], Keys) ->
    SrcKeys2 = Keys -- [Key],
    SrcKeys3 = find_missing(SubTree, SrcKeys2),
    find_missing(RestTree, SrcKeys3).


get_all_key_paths_rev([], KeyPathAcc) ->
    KeyPathAcc;
get_all_key_paths_rev([{Key, Value, SubTree} | RestTree], KeyPathAcc) ->
    get_all_key_paths_rev(SubTree, [{Key, Value} | KeyPathAcc]) ++
        get_all_key_paths_rev(RestTree, KeyPathAcc).
        
    
% Removes any branches from the tree whose leaf node(s) are in the Keys
remove_leafs(Tree, Keys) ->
    % flatten each branch in a tree into a tree path
    Paths = get_all_key_paths_rev(Tree, []),
    
    % filter out any that are in the keys list.
    {FoundKeys, FilteredPaths} = lists:mapfoldl(
        fun(Key, PathsAcc) ->
            case [Path || [{LeafKey,_}|_]=Path <- PathsAcc, LeafKey /= Key] of
            PathsAcc ->
                {nil, PathsAcc};
            PathsAcc2 ->
                {Key, PathsAcc2}
            end
        end, Paths, Keys),
        
    % convert paths back to trees
    NewTree = lists:foldl(
        fun(Path,TreeAcc) ->
            SingleTree = lists:foldl(
                fun({K,V},NewTreeAcc) -> [{K,V,NewTreeAcc}] end, [], Path),
            merge(TreeAcc, SingleTree)
        end, [], FilteredPaths),
    {NewTree, FoundKeys}.


% get the leafs in the tree matching the keys. The matching key nodes can be
% leafs or an inner nodes. If an inner node, then the leafs for that node
% are returned.
get_key_leafs(Tree, Keys) ->
    get_key_leafs(Tree, Keys, []).
    
get_key_leafs(_Tree, [], _KeyPathAcc) ->
    {[], []};
get_key_leafs([], KeysToGet, _KeyPathAcc) ->
    {[], KeysToGet};
get_key_leafs([{Key, _Value, SubTree}=Tree | RestTree], KeysToGet, KeyPathAcc) ->
    case KeysToGet -- [Key] of
    KeysToGet -> % same list, key not found    
        {LeafsFound, KeysToGet2} = get_key_leafs(SubTree, KeysToGet, [Key | KeyPathAcc]),
        {RestLeafsFound, KeysRemaining} = get_key_leafs(RestTree, KeysToGet2, KeyPathAcc),
        {LeafsFound ++ RestLeafsFound, KeysRemaining};
    KeysToGet2 ->
        LeafsFound = get_all_leafs([Tree], KeyPathAcc),
        LeafKeysFound = [LeafKeyFound || {LeafKeyFound, _, _} <- LeafsFound],
        KeysToGet2 = KeysToGet2 -- LeafKeysFound,
        {RestLeafsFound, KeysRemaining} = get_key_leafs(RestTree, KeysToGet2, KeyPathAcc),
        {LeafsFound ++ RestLeafsFound, KeysRemaining}
    end.

get(Tree, KeysToGet) ->
    {KeyPaths, KeysNotFound} = get_full_key_paths(Tree, KeysToGet),
    FixedResults = [ {Key, Value, [Key0 || {Key0, _} <- Path]} || [{Key, Value}|_] = Path <- KeyPaths],
    {FixedResults, KeysNotFound}.
    
get_full_key_paths(Tree, Keys) ->
    get_full_key_paths(Tree, Keys, []).
    
get_full_key_paths(_Tree, [], _KeyPathAcc) ->
    {[], []};
get_full_key_paths([], KeysToGet, _KeyPathAcc) ->
    {[], KeysToGet};
get_full_key_paths([{KeyId, Value, SubTree} | RestTree], KeysToGet, KeyPathAcc) ->
    KeysToGet2 = KeysToGet -- [KeyId],
    CurrentNodeResult =
    case length(KeysToGet2) == length(KeysToGet) of
    true -> % not in the key list.
        [];
    false -> % this node is the key list. return it
        [[{KeyId, Value} | KeyPathAcc]]
    end,
    {KeysGotten, KeysRemaining} = get_full_key_paths(SubTree, KeysToGet2, [{KeyId, Value} | KeyPathAcc]),
    {KeysGotten2, KeysRemaining2} = get_full_key_paths(RestTree, KeysRemaining, KeyPathAcc),
    {CurrentNodeResult ++ KeysGotten ++ KeysGotten2, KeysRemaining2}.

get_all_leafs_full(Tree) ->
    get_all_leafs_full(Tree, []).
    
get_all_leafs_full([], _KeyPathAcc) ->
    [];
get_all_leafs_full([{KeyId, Value, []} | RestTree], KeyPathAcc) ->
    [[{KeyId, Value} | KeyPathAcc] | get_all_leafs_full(RestTree, KeyPathAcc)];
get_all_leafs_full([{KeyId, Value, SubTree} | RestTree], KeyPathAcc) ->
    get_all_leafs_full(SubTree, [{KeyId, Value} | KeyPathAcc]) ++ get_all_leafs_full(RestTree, KeyPathAcc).

get_all_leafs(Tree) ->
    get_all_leafs(Tree, []).
    
get_all_leafs([], _KeyPathAcc) ->
    [];
get_all_leafs([{KeyId, Value, []} | RestTree], KeyPathAcc) ->
    [{KeyId, Value, [KeyId | KeyPathAcc]} | get_all_leafs(RestTree, KeyPathAcc)];
get_all_leafs([{KeyId, _Value, SubTree} | RestTree], KeyPathAcc) ->
    get_all_leafs(SubTree, [KeyId | KeyPathAcc]) ++ get_all_leafs(RestTree, KeyPathAcc).

get_leaf_keys([]) ->
    [];
get_leaf_keys([{Key, _Value, []} | RestTree]) ->
    [Key | get_leaf_keys(RestTree)];
get_leaf_keys([{_Key, _Value, SubTree} | RestTree]) ->
    get_leaf_keys(SubTree) ++ get_leaf_keys(RestTree).
    
count_leafs([]) ->
    0;
count_leafs([{_Key, _Value, []} | RestTree]) ->
    1 + count_leafs(RestTree);
count_leafs([{_Key, _Value, SubTree} | RestTree]) ->
    count_leafs(SubTree) + count_leafs(RestTree).
    

map(_Fun, []) ->
    [];
map(Fun, [{Key, Value, SubTree} | RestTree]) ->
    Value2 = Fun(Key, Value),
    [{Key, Value2, map(Fun, SubTree)} | map(Fun, RestTree)].