#pragma once

#include "RawTree.h"

#include "HuffmanCodeLengthTable.h"
#include "HuffmanFixedCodes.h"

#include <vector>
#include <tuple>
#include <unordered_map>

class PrefixCodeGenerator
{
public:
    virtual ~PrefixCodeGenerator() = default;
    virtual std::optional<PrefixCode> getLiteralValue(unsigned char symbol) const = 0;
    virtual std::optional<PrefixCode> getLengthValue(unsigned length) const = 0;
    virtual std::optional<PrefixCode> getDistanceValue(unsigned distance) const = 0;

    virtual std::optional<PrefixCode> getEndOfStreamValue() const = 0;
};

class HuffmanEncoder : public PrefixCodeGenerator
{
    using CountPair = std::pair<unsigned, unsigned>;
    using Hit = std::tuple<unsigned, unsigned, unsigned char>;

public:
    void encode(const std::vector<unsigned>& counts);
    void encode(const std::unordered_map<unsigned char, unsigned>& counts);

    uint32_t getLengthValue(unsigned length);

    std::optional<PrefixCode> getLiteralValue(unsigned char symbol) const override;

    std::optional<PrefixCode> getLengthValue(unsigned length) const override;

    std::optional<PrefixCode> getDistanceValue(unsigned distance) const override;

    std::optional<PrefixCode> getEndOfStreamValue() const override;

    void initializeTrees(const std::vector<Hit>& hits);

    void setUseFixedCode(bool useFixed);
private:
    void initializeLiteralLengthTable(const std::vector<Hit>& hits);

    void dumpTree(const RawTree<CountPair>& tree) const;
    void dumpNode(RawNode<CountPair>* node, unsigned depth) const;

    bool mUseFixedCode{false};
    bool mTableIsInitialized{false};

    std::vector<unsigned char> mSymbolMapping;
    HuffmanCodeLengthTable mLiteralLengthTable;
    HuffmanCodeLengthTable mDistanceTable;
};