#include "HuffmanTree.h"

#include "ByteUtils.h"

#include <sstream>
#include <algorithm>
#include <iostream>

PrefixCode::PrefixCode(uint32_t data, unsigned length)
    : mLength(length)
{
    mData = ByteUtils::mirror(data, length);
}

bool PrefixCode::matches(unsigned length, uint32_t code) const
{
    return (mLength == length) && (mData == code);
}

std::string PrefixCode::toString(bool bitsAsRightToLeft) const
{
    if (bitsAsRightToLeft)
    {
        if (mLength <=8 )
        {
            return ByteUtils::toString(mData).substr(8 - mLength, mLength);
        }
        else
        {
            return ByteUtils::toString(mData, mLength);
        }
    }
    else
    {
        if (mLength <=8 )
        {
            return ByteUtils::toString(ByteUtils::mirror(mData, mLength)).substr(0, mLength);
        }
        else
        {
            return ByteUtils::toString(mData, mLength);
        }
    }
}

void HuffmanTree::addCodeLengthEntry(unsigned length, const CodeSymbolPair& data)
{
    bool found{false};
    for (auto& entry : mTable)
    {
        if (entry.first == length)
        {
            entry.second.push_back(data);
            found = true;
            break;
        }
    }

    if (!found)
    {
        mTable.push_back({length, {data}});
    }
}

void HuffmanTree::sortTable()
{
    std::sort(mTable.begin(), mTable.end(), [](CodeLengthData a, CodeLengthData b){return a.first < b.first;});
}

std::optional<HuffmanTree::Symbol> HuffmanTree::findMatch(std::size_t treeIndex, uint32_t code) const
{
    const auto& legth_data = mTable[treeIndex];
    for(const auto& entry : legth_data.second)
    {
        //std::cout << "Checking if " << entry.second << " matches code " << ByteUtils::toString(code) << std::endl;;
        if (entry.first.matches(legth_data.first, code))
        {
            return entry.second;
        }
    }
    return std::nullopt;
}

std::optional<PrefixCode> HuffmanTree::getCode(Symbol symbol) const
{
    for(const auto& entry : mTable)
    {
        for(const auto& data : entry.second)
        {
            if (data.second == symbol)
            {
                return data.first;
            }
        }
    }
    return std::nullopt;
}

std::size_t HuffmanTree::getNumCodeLengths() const
{
    return mTable.size();
}

unsigned HuffmanTree::getCodeLength(std::size_t idx) const
{
    return mTable[idx].first;
}

std::string HuffmanTree::dump(bool bitsAsRightToLeft) const
{
    std::stringstream sstr;
    for (const auto& code_length_data : mTable)
    {
        sstr << "Prefix table for Code Length " << code_length_data.first << " has vals: \n";
        for (const auto& entry : code_length_data.second)
        {
            sstr << "Code " << entry.first.toString(bitsAsRightToLeft) << " Symbol: " << entry.second << '\n';
        }
    }
    return sstr.str();
}
