#pragma once

#include <vector>
#include <string>
#include <optional>

class PrefixCode
{
public:
    PrefixCode(uint32_t data, unsigned length);

    std::string toString(bool bitsAsRightToLeft = true) const;

    bool matches(unsigned length, uint32_t code) const;

    uint32_t getData() const
    {
        return mData;
    }

    unsigned getLength() const
    {
        return mLength;
    }

private:
    unsigned mLength{0};
    uint32_t mData{0};
};

class HuffmanTree
{
public:
    using Symbol = unsigned;
    using CodeLength = unsigned;
    using CodeSymbolPair = std::pair<PrefixCode, Symbol>;
    using CodeLengthData = std::pair<CodeLength, std::vector<CodeSymbolPair> >;

    void addCodeLengthEntry(unsigned length, const CodeSymbolPair& data);

    std::string dump(bool bitsAsRightToLeft = true) const;

    std::optional<HuffmanTree::Symbol> findMatch(std::size_t treeIndex, uint32_t code) const;

    std::size_t getNumCodeLengths() const;

    unsigned getCodeLength(std::size_t idx) const;

    std::optional<PrefixCode> getCode(Symbol symbol) const;

    void sortTable();
private:
    std::vector<CodeLengthData> mTable;
};
