#include "ZlibEncoder.h"

#include "ByteUtils.h"
#include "DeflateEncoder.h"
#include "FileLogger.h"
#include "BitStream.h"

#include "Adler32Checksum.h"

#include <math.h>
#include <iostream>
#include <sstream>

ZlibEncoder::ZlibEncoder(BitStream* inputStream, BitStream* outputStream)
: AbstractEncoder(inputStream, outputStream)
{
    mChecksumCalculator = std::make_unique<Adler32Checksum>();
}

ZlibEncoder::~ZlibEncoder()
{

}

void ZlibEncoder::setWindowSize(unsigned size)
{
    mWindowSize = size;
}

std::string ZlibEncoder::toString(CompressionLevel level) const
{
    switch(level)
    {
    case CompressionLevel::FASTEST:
        return "FASTEST";
    case CompressionLevel::FAST:
        return "FAST";
    case CompressionLevel::DEFAULT:
        return "DEFAULT";
    case CompressionLevel::MAX_COMPRESSION:
        return "MAX_COMPRESSION";
    default:
        return "UNKNOWN";
    }
}

std::string ZlibEncoder::toString(CompressionMethod method) const
{
    return method == CompressionMethod::DEFLATE ? "DEFLATE" : "UNKNOWN";
}

void ZlibEncoder::parseCompressionMethod(unsigned char method)
{
    //std::cout << "Got compression input " << static_cast<int>(method) << std::endl;
    mCompressionMethod = static_cast<CompressionMethod>(ByteUtils::getLowerNBits(method, 4));
    auto compression_info = ByteUtils::getHigherNBits(method, 4);

    if (mCompressionMethod == CompressionMethod::DEFLATE)
    {
        mWindowSize = static_cast<unsigned>(pow(2, compression_info + 8));
    }
}

void ZlibEncoder::parseExtraFlags(unsigned char extraFlags, unsigned char compression_byte)
{
    //std::cout << "Got flags  " << static_cast<int>(extraFlags) << std::endl;

    auto mod = ((static_cast<unsigned>(compression_byte) << 8) | extraFlags) % 31;
    if (mod != 0)
    {
        //std::cout << "Invalid header. Mod is " << mod << std::endl;
    }

    mFlagCheck = ByteUtils::getLowerNBits(extraFlags, 5);
    mUseDictionary = bool(ByteUtils::getBitN(extraFlags, 5));
    mFlagLevel = static_cast<CompressionLevel>(ByteUtils::getHigherNBits(extraFlags, 2));
}

std::string ZlibEncoder::getData() const
{
    std::stringstream sstream;
    sstream << "ZlibEncoder data \n";
    sstream << "Compression method: " << toString(mCompressionMethod) << '\n';
    sstream << "Window size: " << mWindowSize << '\n';
    sstream << "Flag check: " << static_cast<int>(mFlagCheck) << '\n';
    sstream << "Use dictionary: " << mUseDictionary << '\n';
    sstream << "Flag level: " << toString(mFlagLevel) << '\n';
    return sstream.str();
}

bool ZlibEncoder::encode()
{
    DeflateEncoder* deflate_encoder{nullptr};
    if (!mWorkingEncoder)
    {
        if (mCompressionMethod == CompressionMethod::DEFLATE)
        {
            auto uq_deflate_encoder = std::make_unique<DeflateEncoder>(mInputStream, mOutputStream);
            deflate_encoder = uq_deflate_encoder.get();
            mWorkingEncoder = std::move(uq_deflate_encoder);
            mWorkingEncoder->addChecksumCalculator(mChecksumCalculator.get());
        }
        else
        {
            MLOG_ERROR("Zib requested decoder not recognized: " << static_cast<int>(mCompressionMethod) << " aborting encode");
            return false;
        }
    }
    deflate_encoder->setCompressionMethod(mDeflateCompressionMethod);

    auto compression_info = static_cast<unsigned char>(log2(mWindowSize) - 8);
    const unsigned char compression_byte = (compression_info << 4) | static_cast<unsigned char>(mCompressionMethod);

    //std::cout << "ZlibEncoder Writing compression byte " << static_cast<int>(compression_byte) << " with info " << static_cast<int>(compression_info) << std::endl;
    mOutputStream->writeByte(compression_byte);

    unsigned char flag_byte{0};
    flag_byte |= (static_cast<unsigned char>(mUseDictionary) << 5);
    flag_byte |= (static_cast<unsigned char>(mFlagLevel) << 6);

    const auto mod = (unsigned(compression_byte)*256 + flag_byte) % 31;
    flag_byte += (31 - mod);

    //std::cout << "ZlibEncoder Writing Flag byte " << static_cast<int>(flag_byte) << std::endl;
    mOutputStream->writeByte(flag_byte);

    if(!mWorkingEncoder->encode())
    {
        MLOG_ERROR("Sub-Encoder failed - aborting zlib encode");
        //return false;
    }

    const auto checksum = mChecksumCalculator->getChecksum();
    //std::cout << "ZlibEncoder Writing Adler32 Checksum " << checksum << std::endl;
    mOutputStream->write(checksum);
    return true;
}

bool ZlibEncoder::decode()
{
    auto compression_byte = *mInputStream->readNextByte();

    parseCompressionMethod(compression_byte);
    parseExtraFlags(*mInputStream->readNextByte(), compression_byte);

    if (!mWorkingEncoder)
    {
        if (mCompressionMethod == CompressionMethod::DEFLATE)
        {
            mWorkingEncoder = std::make_unique<DeflateEncoder>(mInputStream, mOutputStream);
        }
        else
        {
            MLOG_ERROR("Zib requested decoder not recognized: " << static_cast<int>(mCompressionMethod) << " aborting decode");
            return false;
        }
    }

    auto valid = mWorkingEncoder->decode();

    unsigned char byte0 = *mInputStream->readNextByte();
    unsigned char byte1 = *mInputStream->readNextByte();
    unsigned char byte2 = *mInputStream->readNextByte();
    unsigned char byte3 = *mInputStream->readNextByte();

    uint32_t adler32 = (byte0 << 24) | (byte1 << 16) | (byte2 << 8) | byte3;
    (void)adler32;
    //std::cout << "Got adler 32 checksum " << adler32 << std::endl;

    return valid;
}
