-module (bencode).
-export ([encode/1, decode/1]).

%% Simple bencoding definition:
%% byte strings: <ascii number in base 10>:contents
%% integers: i<ascii number in base 10>e
%% lists: l<contents>e
%% dictionaries: d<contents>e, contents are string keyed in lexographic order

%% ============
%% = Exported =
%% ============

%% encode a Term
encode({dictionary, Entries}) when is_list(Entries) ->
    SortedEntries = lists:keysort(1, Entries),
    EncodedEntries = encode_dictionary_entries(SortedEntries),
    <<$d, EncodedEntries/bytes, $e>>;
encode(Data) when is_list(Data) ->
    EncodedList = encode_list(Data),
    <<$l, EncodedList/bytes, $e>>;
encode(Data) when is_binary(Data) ->
    SizeByteString = integer_to_binary_string(size(Data)),
    <<SizeByteString/binary, $:, Data/binary>>;
encode(Data) when is_integer(Data) ->
    Bin = integer_to_binary_string(Data),
    <<$i, Bin/binary, $e>>.

%% decode a Term
decode(<<$d, Rest/binary>>)->
    {dictionary, decode_dictionary(Rest)};
decode(<<$l, Rest/bits>>) ->
    decode_list(Rest);
decode(<<$i, Rest/bits>>) ->
    String = extract_list_from_binary(Rest, $e),
    list_to_integer(String);
decode(<<Byte, Rest/bits>>) when Byte >= $0, Byte =< $9 ->
    Bin = <<Byte, Rest/bits>>,
    SizeString = extract_list_from_binary(Bin, $:),
    Offset = length(SizeString) * 8,
    Size = list_to_integer(SizeString),
    BitsToRead = Size * 8,
    try
	<<_SizeBytes:Offset, $:, BString:BitsToRead/bits, _MoreBin/bits>> = Bin,
	BString
    catch
	error:{badmatch, Bin} ->
	    io:format("error in bencoded integer near ~p~n", [binary_to_list(Bin)])
    end.

%% ===========
%% = Helpers =
%% ===========

%% decode a dictionary
decode_dictionary(<<$e, _Rest/bits>>) -> [];
decode_dictionary(Bin) ->
    {Key, Rest1} = unshift_term(Bin),
    {Value, Rest2} = unshift_term(Rest1),
    [{Key, Value}|decode_dictionary(Rest2)].

%% decode a list
decode_list(<<$e, _Rest/bits>>) -> [];
decode_list(Bin) ->
    {Term, Rest} = unshift_term(Bin),
    [Term|decode_list(Rest)].

%% decode one term from a bencoded binary string
unshift_term(Bin) ->
    Term = decode(Bin),
    EncodedTerm = encode(Term),
    Size = size(EncodedTerm),
    <<EncodedTerm:Size/bytes, Rest/bits>> = Bin,
    {Term, Rest}.

%% extract a list of bytes from a binary stopping at Delimiter
extract_list_from_binary(<<Byte, Delimiter, _Rest/bits>>, Delimiter) -> [Byte];
extract_list_from_binary(<<Byte, Rest/bits>>, Delimiter) ->
    [Byte|extract_list_from_binary(Rest, Delimiter)].

%% return a binary string representation of an integer in binary form
integer_to_binary_string(N) -> list_to_binary(integer_to_list(N)).

%% encode a list of terms
encode_list([H|T]) ->
    EncodedHead = encode(H),
    EncodedTail = encode_list(T),
    <<EncodedHead/binary, EncodedTail/binary>>;
encode_list([]) -> <<>>.

%% encode a list of of {binary, Data}
encode_dictionary_entries([{Name, Value}|Rest]) ->
    EncodedName = encode(Name),
    EncodedValue = encode(Value),
    EncodedRest = encode_dictionary_entries(Rest),
    <<EncodedName/binary, EncodedValue/binary, EncodedRest/binary>>;
encode_dictionary_entries([]) -> <<>>.