#include "XMemDecompress.h"

#include "Utils/Logging/Log.h"

#include <format>
#include <lzx.h>

namespace
{
    uint8_t NextByte(const uint8_t* input, size_t& offset, size_t& remainingSize)
    {
        const auto value = input[offset];
        offset++;
        remainingSize--;
        return value;
    }

    uint16_t CombineHighLow(const uint8_t highByte, const uint8_t lowByte)
    {
        return static_cast<uint16_t>(static_cast<uint16_t>(static_cast<uint16_t>(highByte) << 8u) | static_cast<uint16_t>(lowByte));
    }

    void LogErrorHeaderSpace(size_t remainingInputSize)
    {
        con::error("XMemDecompress: Not enough data for header: {}", remainingInputSize);
    }
} // namespace

XMemDecompressContext::XMemDecompressContext()
    : m_lzx_state(lzx_init(17))
{
}

XMemDecompressContext::~XMemDecompressContext()
{
    if (m_lzx_state)
        lzx_teardown(static_cast<lzx_state*>(m_lzx_state));
}

XMemDecompressContext::XMemDecompressContext(XMemDecompressContext&& other) noexcept
    : m_lzx_state(other.m_lzx_state)
{
    other.m_lzx_state = nullptr;
}

XMemDecompressContext& XMemDecompressContext::operator=(XMemDecompressContext&& other) noexcept
{
    m_lzx_state = other.m_lzx_state;
    other.m_lzx_state = nullptr;
    return *this;
}

void XMemDecompressContext::Reset() const
{
    lzx_reset(static_cast<lzx_state*>(m_lzx_state));
}

std::optional<size_t> XMemDecompressContext::Process(const uint8_t* input, const size_t inputLength, uint8_t* output, const size_t outputBufferSize) const
{
    size_t curInputOffset = 0uz;
    size_t curInputSize = inputLength;

    size_t curOutputOffset = 0uz;
    size_t curOutputSize = outputBufferSize;

    uint8_t lowByte;
    uint16_t dstSize, srcSize;

    while (curInputSize > 0)
    {
        uint8_t highByte = NextByte(input, curInputOffset, curInputSize);

        uint8_t suffixSize;
        if (highByte == 0xFF) // magic number: output is smaller than 0x8000
        {
            if (curInputSize < 4)
            {
                LogErrorHeaderSpace(curInputSize);
                return std::nullopt;
            }

            highByte = NextByte(input, curInputOffset, curInputSize);
            lowByte = NextByte(input, curInputOffset, curInputSize);
            dstSize = CombineHighLow(highByte, lowByte);

            highByte = NextByte(input, curInputOffset, curInputSize);
            lowByte = NextByte(input, curInputOffset, curInputSize);
            srcSize = CombineHighLow(highByte, lowByte);

            // The game seems to skip a 5 byte suffix after these blocks, not sure why.
            suffixSize = 5u;
        }
        else
        {
            if (curInputSize < 1)
            {
                LogErrorHeaderSpace(curInputSize);
                return std::nullopt;
            }

            dstSize = 0x8000u;
            lowByte = NextByte(input, curInputOffset, curInputSize);
            srcSize = CombineHighLow(highByte, lowByte);

            suffixSize = 0u;
        }

        if (srcSize == 0 || dstSize == 0)
        {
            // Other implementations do not handle this as a failure, game code suggests otherwise though
            con::error("XMemDecompress: EOF: {} {}, {}", srcSize, dstSize, curInputSize);
            return std::nullopt;
        }

        if (static_cast<size_t>(srcSize) + suffixSize > curInputSize)
        {
            con::error("XMemDecompress: block size bigger than remaining data: {} > {}", srcSize, curInputSize);
            return std::nullopt;
        }

        if (dstSize > curOutputSize)
        {
            con::error("XMemDecompress: output size bigger than remaining data: {} > {}", dstSize, curOutputSize);
            return std::nullopt;
        }

        auto ret = lzx_decompress(static_cast<lzx_state*>(m_lzx_state), &input[curInputOffset], &output[curOutputOffset], srcSize, dstSize);
        curInputOffset += srcSize + suffixSize;
        curInputSize -= (srcSize + suffixSize);
        curOutputOffset += dstSize;
        curOutputSize -= srcSize;

        if (ret != DECR_OK)
        {
            con::error("XMemDecompress: lzx decompression failed: {}", ret);
            return std::nullopt;
        }
    }

    return curOutputOffset;
}
