#pragma once

#include "RawTree.h"

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

#include <vector>
#include <unordered_map>

class PrefixCodeGenerator
{
public:
    virtual ~PrefixCodeGenerator() = default;
    virtual const PrefixCode& getLiteralValue(unsigned char value) const = 0;

    virtual const PrefixCode& getEndOfStreamValue() const = 0;

};

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

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

    void setUseFixedCode(bool useFixed)
    {
        mUseFixedCode = useFixed;
    }

    uint32_t getLengthValue(unsigned length)
    {
        return 0;
    }

    const PrefixCode& getLiteralValue(unsigned char value) const override
    {
        return mLiteralLengthTable.getCode(value);
    }

    const PrefixCode& getEndOfStreamValue() const override
    {
        return mLiteralLengthTable.getCode(256);
    }

    void initializeLiteralLengthTable()
    {
        if(mUseFixedCode)
        {
            mLiteralLengthTable.setInputLengthSequence(HuffmanFixedCodes::getDeflateFixedHuffmanCodes(), false);
            mLiteralLengthTable.buildPrefixCodes();
        }
    }

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

    bool mUseFixedCode{false};
    bool mTableIsInitialized{false};
    HuffmanCodeLengthTable mLiteralLengthTable;
    HuffmanCodeLengthTable mDistanceTable;
};