divingbells

%%% @author Robert Lasu
%%% @copyright 2025 Robert Lasu <robert.lasu@gmail.com>
%%% @doc Huffman code tree builder, bit encoder and header generator.
%%% 
%%% Module to handle all huffman coding. Compress or decompress a bitstring, can be configured through
%%% a list of tuples.  
%%% @end

-module(huffman).

-ifdef(TEST).
-compile(export_all).
-endif.

-export(['_decompress'/2, '_compress'/2]).

-define(ERR(Msg), {error, {?MODULE, ?LINE, Msg}}).
-define(TBI(Msg), {tbi, {?MODULE, ?LINE, Msg}}).



-record(huffNode, {left, right, freq}).
-opaque huffNode() :: #huffNode{left :: huffNode(), right :: huffNode(), freq :: integer()}.

-record(leaf, {value, freq}).
-opaque leaf() :: #leaf{value :: integer(), freq :: integer()}.

-record(header, {tree, bitExtension}).
-opaque header() :: #header{tree :: huffNode(), bitExtension :: integer()}.



-spec new_tree(Bits) -> {ok, Tree} | {error, Reason} when
    Bits    :: bitstring(),
    Tree    :: huffNode(),
    Reason  :: term().
%%% @doc             Creates a new huffman tree based on a bitstring.
%%% 
%%% `Bits' is a bitstring representation of a file.
%%% `new_tree/1' returns `{ok, Tree}' where Tree is a huffNode().
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
new_tree(Bits) ->
    case get_sorted_frequency_list(Bits) of
        {ok, SortedList} ->
            {ok, create_tree(SortedList)};
        Default ->
            Default
    end.


-spec create_tree(List) -> CompleteTree when
    List            :: list(),
    CompleteTree    :: huffNode().
%%% @doc             Creates a new huffman tree, based on a list of tuples.
%%% 
%%% `List' contains tuples of a binary and the frequency of that binary.
%%% `create_tree/1' returns `{ok, Tree}' where Tree is a huffnode().
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
create_tree([{_, CompleteTree} | []]) -> CompleteTree;
create_tree([{Freq1, Node1} | [{Freq2, Node2} | SortedList]]) ->
    Tree = #huffNode{left = Node1, right = Node2, freq = Freq1+Freq2},
    create_tree(lists:sort([{Freq1+Freq2, Tree} | SortedList])).


-spec get_header(Bits) -> {ok, {BencodeHeader, EncodedBits}} | {error, Reason} when
    Bits            :: bitstring(),
    BencodeHeader   :: string(),
    EncodedBits     :: bitstring(),
    Reason          :: term().
%%% @doc             Get the header of a encoded bitstring.
%%%
%%% `Bits' is a bitstring of encoded data, which the header will be extracted from.
%%% Returns `{ok, {Header, EncodedBits}}' where `Header' is a header() for `EncodedBits'.
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end 
%%% @since           0.0.1
get_header(Bits) ->
    case get_header_size(Bits, []) of
        {ok, {HeaderSize, HeaderAndBits}} ->
            case get_header(HeaderAndBits, HeaderSize, []) of
                {ok, {Header, EncodedBits}} ->
                    {ok, BencodeHeader} = bencode:decode(Header),
                    {ok, {BencodeHeader, EncodedBits}};
                Default ->
                    Default
            end;
        Default ->
            Default
    end.


-spec get_header(CompressedBits, HeaderSize, HeaderList) -> {ok, {Header, Payload}} | {error, Reason} when
    CompressedBits      :: bitstring(),
    HeaderSize          :: integer(),
    HeaderList          :: list(),
    Header              :: string(),
    Payload             :: bitstring(),
    Reason              :: term().
%%% @doc             Extract the header from some compressed bits.
%%% 
%%% Append the head binary of `CompressedBits' to `HeaderList' for `HeaderSize' number of times.
%%% Returns `{ok, {Header, Payload}}' where `Header' is a the header bits, and `Payload' is the compressed data.
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
get_header(CompressedBits, 0, List) -> {ok, {lists:reverse(List), CompressedBits}};
get_header(<<Char:8, CompressedBits/bitstring>>, HeaderSize, List) -> get_header(CompressedBits, HeaderSize-1, [Char | List]);
get_header(_,_,_) -> ?ERR(no_match).


-spec get_header_size(Bits, List) -> {ok, {Num, Rest}} | {error, Reason} when
    Bits    :: bitstring(),
    List    :: list(),
    Num     :: integer(),
    Rest    :: bitstring(),
    Reason  :: term().
%%% @doc             Get the header size.
%%%
%%% Append the header binary, `N', of `Bits' to `List' iff `N' is a string representation of a number, until the char ":" is read.
%%% Returns `{ok, {Num, Rest}}' where `Num' is the length of the header and `Rest' is the remaining data.
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
get_header_size(<<58:8, Bits/bitstring>>, List) -> {ok, {list_to_integer(lists:reverse(List)), Bits}};
get_header_size(<<Num:8, Bits/bitstring>>, List) when Num >= 48, Num =< 57 -> get_header_size(Bits, [Num | List]);
get_header_size(_,_) -> ?ERR(malformed_header).


-spec create_header(Tree, Compressed) -> {header, {HeaderSize, Header, Extend}} | {error, Reason} when
    Tree        :: huffNode(),
    Compressed  :: bitstring(),
    HeaderSize  :: bitstring(),
    Header      :: bitstring(),
    Extend      :: integer(),
    Reason      :: term().
                            %%% {ok, HeaderBinBen, BitExtension} which already creates the header, to-be, record
                            %%% and transform it to bencode.
                            %%% HeaderBinBen    = <<HeaderBencodeSize, ":", HeaderBencode>>
                            %%% HeaderBencode    = Bencode of {header, {Tree, BitExtension}}
                            %%% Tree             = huffNode()
                            %%% BitExtension     = integer()
%%% @doc             Creates a header for some compressed data.
%%%
%%% Returns a header `{header, {HeaderSize, Header, Extended}}' for the bits `Compressed'.
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
create_header(Tree, Compressed) ->
    Extend = (8 - (bit_size(Compressed) rem 8)) rem 8,
    {ok, HeaderBencode} = bencode:encode(#header{tree = Tree, bitExtension = Extend}),
    HeaderBincode = list_to_binary(HeaderBencode),
    HeaderBencodeSize = list_to_binary(integer_to_list(string:len(HeaderBencode))),
    HeaderBinBen = <<HeaderBencodeSize/bitstring, ":":8, HeaderBincode/bitstring>>,
    {ok, HeaderBinBen, Extend}.


-spec serialise_tree(Tree) -> {ok, Encodede} |  {error, Reason} when
    Tree        :: huffNode(),
    Encodede    :: string(),
    Reason      :: term().
%%% @doc             Serialise a huffNode().
%%%
%%% Serilise the huffNode() record, `Tree', to bencode in order to write the data structure to file.
%%% Returns `{ok, Encoded}' where `Encoded' is the bencode representation of `Tree'.
%%% Throws `{error, Reason}' if any errors should arise.
%%% @end
%%% @since           0.0.1
serialise_tree(Tree) -> bencode:encode(Tree).


-spec get_sorted_frequency_list(Bits) -> {ok, List} | {error, Reason} when
    Bits    :: bitstring(),
    List    :: list(),
    Reason  :: term().
%%% @doc             Returns a sorted list of binaries frequency.
%%%
%%% Read every byte, back-to-back, of `Bits' and calculate the frequency the occurrencies of that binary.
%%% The window is 8 bits, non-sliding window.
%%% Returns `{ok, List}' where `List' is a list of tuple `{Freq, Leaf}' 
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @since           0.0.1
get_sorted_frequency_list(Bits) -> 
    case lists:sort(get_frequency_list(Bits)) of
        Err = {error, _} ->
            Err;
        Default ->
            {ok, Default}
    end.


-spec get_frequency_list(Bits) -> [{Freq, Leaf}] | {error, Reason} when
    Bits    :: bitstring(),
    Freq    :: integer(),
    Leaf    :: leaf(),
    Reason  :: term().
%%% @doc             Get the frequency of each back-to-back binary
%%% 
%%% Read every byte, back-to-back, of `Bits' and calculate the frequency of that binary.
%%% The window is 8 bits, non-sliding window.
%%% Returns a list of tuples `{Freq, Leaf}'
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @since           0.0.1
get_frequency_list(<<>>) -> [];
get_frequency_list(<<Char:8, ByteArray/bitstring>>) ->
    case get_frequency(Char, ByteArray) of
        {ok, {Freq, Rest}} ->
            [{Freq, #leaf{value = Char, freq = Freq}} | get_frequency_list(Rest)];
        Default ->
            Default
    end.


-spec get_frequency(Char, Bits) -> {ok, {Frequency, Rest}} | {error, Reason} when
    Char        :: string(),
    Bits        :: bitstring(),
    Frequency   :: integer(),
    Rest        :: bitstring(),
    Reason      :: term().
%%% @doc             Get the frequency of a char in a bitstring
%%% 
%%% Calculates the frequency of `Char' in `Bits'.
%%% Returns `{ok, {Frequecy, Rest}}' where `Frequency' is the number of occurrencies of `Char' + 1.
%%% `Rest' is the remaining bits without any occurrencies of `Char'.
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see get_frequency/3
%%% @since           0.0.1
get_frequency(Char, Bits) -> get_frequency(Char, Bits, {1, <<>>}).


-spec get_frequency(Char, Bits, Data) -> {ok, {Frequency, Rest}} | {error, Reason} when
    Char        :: string(),
    Bits        :: bitstring(),
    Data        :: tuple(),
    Frequency   :: integer(),
    Rest        :: bitstring(),
    Reason      :: term().
%%% @doc             Get the frequency of a char in a bitstring
%%% 
%%% Calculates the frequency of `Char' in `Bits'.
%%% Returns `{ok, {Frequecy, Rest}}' where `Frequency' is the number of occurrencies of `Char' + 1.
%%% `Rest' is the remaining bits without any occurrencies of `Char'.
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see get_frequency/2
%%% @since           0.0.1
get_frequency(_, <<>>, Data) -> {ok, Data};
get_frequency(Char, <<Char, Bits/bitstring>>, {Num, NewBits})    -> get_frequency(Char, Bits, {Num+1, NewBits});
get_frequency(Char, <<NotChar, Bits/bitstring>>, {Num, NewBits}) -> get_frequency(Char, Bits, {Num, <<NotChar, NewBits/binary>>});
get_frequency(_,_,_) -> ?TBI(no_match).


-spec get_code(Tree, Char) -> {ok, Code} | {error, Reason} when
    Tree    :: huffNode() | leaf(),
    Char    :: string(),
    Code    :: bitstring(),
    Reason  :: term().
%%% @doc             Get the huffman code representation of a binary.
%%% 
%%% Returns `{ok, Code}', where `Code' is the huffman code representation of `Char' when encoding with huffNode() `Tree'.
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @since           0.0.1
get_code(Leaf, Char) when is_record(Leaf, leaf), Leaf#leaf.value =:= Char -> {ok, <<>>};
get_code(Leaf, _) when is_record(Leaf, leaf) -> {error, not_in_tree};
get_code(HuffmanTree, Char) when is_record(HuffmanTree, huffNode) ->
    case get_code(HuffmanTree#huffNode.left, Char) of
        {ok, Code} ->
            {ok, <<0:1, Code/bitstring>>};
        _ ->
            case get_code(HuffmanTree#huffNode.right, Char) of
                {ok, Code} ->
                    {ok, <<1:1, Code/bitstring>>};
                Default ->
                    Default
            end
    end.


-spec '_compress'(Bits, Config) -> {ok, CompressedBits} | {error, Reason} when
    Bits            :: bitstring(),
    Config          :: list(),
    CompressedBits  :: bitstring(),
    Reason          :: term().
%%% // TODO Config
%%% @doc             Compress some arbitrary bits.
%%% 
%%% Compress `Bits' with configuration (WIP) `Config'
%%% Return `{ok, CompressedBits}'
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see compress/2
%%% @since           0.0.1
'_compress'(Bits, _Config) -> 
    case new_tree(Bits) of
        {ok, Tree} ->
            case compress(Tree, Bits) of
                {ok, Code} ->
                    case create_header(Tree, Code) of
                        {ok, HeaderBinBen, BitExtension} ->
                            {ok, <<HeaderBinBen/binary, Code/bitstring, 0:BitExtension>>};
                        Default ->
                            Default
                    end;
                Default ->
                    Default
            end;
        Default ->
            Default
    end.


-spec compress(Tree, Bits) -> {ok, Code} | {error, Reason} when
    Tree    :: huffNode(),
    Bits    :: bitstring(),
    Code    :: bitstring(),
    Reason  :: term().
%%% @doc             Compress some arbitrary bits.
%%% 
%%% Compress `Bits' with configuration (WIP) `Config'
%%% Return `{ok, CompressedBits}'
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see '_compress'/2
%%% @since           0.0.1
compress(_, <<>>) -> {ok, <<>>};
compress(HuffmanTree, <<Char, Bitstring/bitstring>>) ->
    case get_code(HuffmanTree, Char) of
        {ok, Code} ->
            case compress(HuffmanTree, Bitstring) of
                {ok, CompressedCode} ->
                    {ok, <<Code/bitstring, CompressedCode/bitstring>>};
                Default ->
                    Default
            end;
        Default ->
            Default
    end;
compress(_,_) -> ?ERR(debugg).


-spec decompress_char(Tree, Bits) -> {ok, {Value, Rest}} | {error, Reason} when
    Tree    :: huffNode() | leaf(),
    Bits    :: bitstring(),
    Value   :: string(),
    Rest    :: bitstring(),
    Reason  :: term().
%%% @doc             Decompress a char.
%%% 
%%% Decompress the char `Bits' using the huffNode() `Tree'.
%%% Return `{ok, {Value, Rest}}' where `Value' is the decompressed char and `Rest' is the remaining compressed bits.
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see '_decompress'/2
%%% @since           0.0.1
decompress_char(Leaf, Rest) when is_record(Leaf, leaf) -> {ok, {Leaf#leaf.value, Rest}};
decompress_char(HuffmanTree, <<>>) when is_record(HuffmanTree, huffNode) -> ?ERR(no_bits_leaf_not_reached);
decompress_char(HuffmanTree, <<0:1, Bitstring/bitstring>>) when is_record(HuffmanTree, huffNode) -> decompress_char(HuffmanTree#huffNode.left, Bitstring);
decompress_char(HuffmanTree, <<1:1, Bitstring/bitstring>>) when is_record(HuffmanTree, huffNode) -> decompress_char(HuffmanTree#huffNode.right, Bitstring).


-spec '_decompress'(Bits, Config) -> {ok, List} | {error, Reason} when
    Bits    :: bitstring(),
    Config  :: list(),
    List    :: list(),
    Reason  :: term().
%%% // TODO Config
%%% @doc             Deompress bits, with some configuration.
%%% 
%%% Decoompress `Bits' with configuration (WIP) `Config'.
%%% Return `{ok, List}' where `List' is the decompressed data. 
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see decompress/2
%%% @see decompress_char/2
%%% @since           0.0.1
'_decompress'(Bits, _Config) ->
    case get_header(Bits) of
        {ok, {Header, CompressedBits}} ->
            Tree = Header#header.tree,
            SizeToDecompress = bit_size(CompressedBits) - Header#header.bitExtension,
            <<ToDecompress:SizeToDecompress, _:(Header#header.bitExtension)>> = CompressedBits,
            BitsToDecompress = <<ToDecompress:SizeToDecompress>>,
            {ok, decompress(Tree, BitsToDecompress)};
            % case bencode:decode(Bencode) of

            %     %%% // TODO Change return value of create header to
            %     %%% {ok, HeaderBinBen, BitExtension} which already creates the header, to-be, record
            %     %%% and transform it to bencode.
            %     %%% HeaderBinBen    = <<HeaderBencodeSize, ":", HeaderBencode>>
            %     %%% HeaderBencode    = Bencode of {header, {Tree, BitExtension}}
            %     %%% Tree             = huffNode()
            %     %%% BitExtension     = integer()


            %     {ok, {header, {Tree, ExtendBits}}} ->
            %         Size = bit_size(CompressedBits) - ExtendBits,
            %         <<ToCompress:Size, _:ExtendBits>> = CompressedBits,
            %         {ok, decompress(Tree, <<ToCompress/binary>>)};
            %     Default ->
            %         Default
            % end;
        Default ->
            Default
    end.


-spec decompress(HuffmanTree, Bits) -> List | {error, Reason} when
    HuffmanTree     :: huffNode(),
    Bits            :: bitstring(),
    List            :: string(),
    Reason          :: term().
%%% @doc             Deompress bits, with some configuration.
%%% 
%%% Decoompress `Bits' with configuration (WIP) `Config'.
%%% Return `List' is the decompressed data. 
%%% Throws `{error, Reason}' if any error aries.
%%% @end
%%% @see '_decompress'/2
%%% @since           0.0.1
decompress(_, <<>>) -> [];
decompress(HuffmanTree, <<Bitstring/bitstring>>) ->
    case decompress_char(HuffmanTree, Bitstring) of
        {ok, {Char, Rest}} ->
            [ Char | decompress(HuffmanTree, Rest) ];
        Default ->
            Default
    end.