#include <string>
#include <memory>
#include <vector>
#include <unordered_map>
#include <cstdint>

class huffman_tree
{
    std::string _value;
    std::shared_ptr<huffman_tree> _left;
    std::shared_ptr<huffman_tree> _right;

public:
    huffman_tree(const std::string& value);
    huffman_tree(std::shared_ptr<huffman_tree> left, std::shared_ptr<huffman_tree> right);

    std::uint8_t value() const;
    std::shared_ptr<huffman_tree> left() const;
    std::shared_ptr<huffman_tree> right() const;

    bool is_leaf() const;
};

struct huffman_probability
{
    std::string bytes;
    double probability;
    std::shared_ptr<huffman_tree> tree;

    huffman_probability(const std::string& bytes, double probability);

    huffman_probability operator+(const huffman_probability& that) const;
    bool operator<(const huffman_probability& that) const;
    bool operator>(const huffman_probability& that) const;
    bool operator<=(const huffman_probability& that) const;
    bool operator>=(const huffman_probability& that) const;
};

using coding = std::vector<bool>;
using probability_list = std::vector<huffman_probability>;
using probability_iterator = probability_list::iterator;
using probability_const_iterator = probability_list::const_iterator;

class huffman_table
{
    std::unordered_map<std::byte, coding> payload;

    void create_table(std::shared_ptr<huffman_tree> tree, coding path);

public:
    huffman_table(std::shared_ptr<huffman_tree> tree);
    huffman_table(std::vector<std::uint8_t>::const_iterator& coding_stream);
    huffman_table() = default;

    void insert(std::uint8_t byte, coding&& path);
    void encode(std::ostream& output) const;
    std::size_t size() const;
    std::unordered_map<std::byte, coding>::iterator begin();
    std::unordered_map<std::byte, coding>::iterator end();
    std::unordered_map<std::byte, coding>::const_iterator begin() const;
    std::unordered_map<std::byte, coding>::const_iterator end() const;

    const coding& operator[](std::uint8_t byte) const;
};

std::shared_ptr<huffman_tree> create_tree(probability_list& probabilities);
probability_list adapt_probabilities(const std::vector<std::uint8_t>& input);
void compress(const std::vector<std::uint8_t>& input, std::ostream& output);
void decompress(const std::vector<std::uint8_t>& input, std::ostream& output);