#include "compression.h"
#include "minilzo.h"
#include "xcompress.h"

#include <QLibrary>
#include <QDebug>
#include <QFile>

QByteArray Compression::CompressXMem(const QByteArray &data)
{
    XMEMCODEC_PARAMETERS_LZX lzxParams = {};
    lzxParams.Flags = 0;
    lzxParams.WindowSize = 0x20000;
    lzxParams.CompressionPartitionSize = 0x80000;

    XMEMCOMPRESSION_CONTEXT ctx = nullptr;
    if (FAILED(XMemCreateCompressionContext(XMEMCODEC_LZX, &lzxParams, 0, &ctx)) || !ctx)
        return QByteArray();

    SIZE_T estimatedSize = data.size() + XCOMPRESS_LZX_BLOCK_GROWTH_SIZE_MAX;
    QByteArray output(static_cast<int>(estimatedSize), 0);
    SIZE_T actualSize = estimatedSize;

    HRESULT hr = XMemCompress(ctx, output.data(), &actualSize, data.constData(), data.size());
    XMemDestroyCompressionContext(ctx);

    if (FAILED(hr))
        return QByteArray();

    output.resize(static_cast<int>(actualSize));
    return output;
}

QByteArray Compression::DecompressXMem(const QByteArray &data,
                                       int flags, int windowSize, int partSize)
{
    if (data.isEmpty())
        return {};

    XMEMCODEC_PARAMETERS_LZX lzxParams = {};
    lzxParams.Flags = flags;
    lzxParams.WindowSize = windowSize;
    lzxParams.CompressionPartitionSize = partSize;

    QByteArray internalState(0x94933, Qt::Uninitialized);

    XMEMDECOMPRESSION_CONTEXT ctx = XMemInitializeDecompressionContext(
        XMEMCODEC_LZX, &lzxParams, 1,
        internalState.data(), internalState.size());

    if (!ctx || XMemResetDecompressionContext(ctx)) {
        qWarning() << "Failed to init LZX context";
        return {};
    }

    QByteArray output;
    output.reserve(16 * 1024 * 1024); // rough guess

    const quint8 *nextIn = reinterpret_cast<const quint8*>(data.constData());
    SIZE_T availIn = data.size();

    QByteArray scratch(0x10000, Qt::Uninitialized); // 64 KB chunks

    while (availIn > 0) {
        SIZE_T inSize = availIn;           // let XMem tell us how much it will consume
        SIZE_T outSize = scratch.size();   // max 64 KB per call

        HRESULT hr = XMemDecompressStream(ctx,
                                          scratch.data(), &outSize,
                                          nextIn, &inSize);

        if (FAILED(hr)) {
            qWarning() << "XMemDecompressStream failed, hr=" << hr;
            XMemDestroyDecompressionContext(ctx);
            return {};
        }

        if (inSize == 0 && outSize == 0)
            break; // no progress

        output.append(scratch.constData(), static_cast<int>(outSize));

        nextIn += inSize;
        availIn -= inSize;
    }

    XMemDestroyDecompressionContext(ctx);
    return output;
}

quint32 Compression::CalculateAdler32Checksum(const QByteArray &data) {
    // Start with the initial value for Adler-32
    quint32 adler = adler32(0L, Z_NULL, 0);

    // Calculate Adler-32 checksum
    adler = adler32(adler, reinterpret_cast<const Bytef *>(data.constData()), data.size());

    return adler;
}

qint64 Compression::FindZlibOffset(const QByteArray &bytes)
{
    QDataStream stream(bytes);

    while (!stream.atEnd())
    {
        QByteArray testSegment = stream.device()->peek(2).toHex().toUpper();
        if (testSegment == "7801" ||
            testSegment == "785E" ||
            testSegment == "789C" ||
            testSegment == "78DA") {
            return stream.device()->pos();
        }
        stream.skipRawData(1);
    }
    return -1;
}

QByteArray Compression::StripHashBlocks(const QByteArray &raw,
                                        int dataChunkSize,
                                        int hashChunkSize)
{
    QByteArray cleaned;
    cleaned.reserve(raw.size());        // upper bound

    int p = 0;
    while (p < raw.size())
    {
        const int chunk = qMin(dataChunkSize, raw.size() - p);
        cleaned.append(raw.constData() + p, chunk);
        p += chunk;

        // skip hash bytes if they are still inside the buffer
        if (p < raw.size())
            p += qMin(hashChunkSize, raw.size() - p);
    }
    return cleaned;
}

QByteArray Compression::DecompressZLIB(const QByteArray &aCompressedData) {
    if (aCompressedData.isEmpty()) {
        return {};
    }

    z_stream strm{};
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    strm.avail_in = static_cast<uInt>(aCompressedData.size());
    strm.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(aCompressedData.data()));

    if (inflateInit2(&strm, MAX_WBITS) != Z_OK) {
        qWarning() << "inflateInit2 failed";
        return {};
    }

    QByteArray decompressed;
    QByteArray buffer(fmin(strm.avail_in * 2, 4096), Qt::Uninitialized);

    int ret;
    do {
        strm.next_out = reinterpret_cast<Bytef*>(buffer.data());
        strm.avail_out = buffer.size();

        ret = inflate(&strm, Z_NO_FLUSH);

        if (strm.avail_out < buffer.size()) {
            decompressed.append(buffer.constData(), buffer.size() - strm.avail_out);
        }

        if (ret == Z_STREAM_END) {
            break;
        }

        if (ret == Z_BUF_ERROR && strm.avail_out == 0) {
            buffer.resize(buffer.size() * 2);
        } else if (ret != Z_OK) {
            size_t errorOffset = strm.total_in;
            qWarning() << "Zlib error:" << zError(ret)
                       << "at offset" << errorOffset
                       << "of" << aCompressedData.size() << "bytes";
            inflateEnd(&strm);
            return decompressed;
        }
    } while (ret != Z_STREAM_END);

    inflateEnd(&strm);
    return decompressed;
}


QByteArray Compression::CompressZLIB(const QByteArray &aData) {
    return CompressZLIBWithSettings(aData);
}

QByteArray Compression::CompressZLIBWithSettings(const QByteArray &aData, int aCompressionLevel, int aWindowBits, int aMemLevel, int aStrategy, const QByteArray &aDictionary) {
    if (aData.isEmpty())
        return {};

    z_stream strm{};
    if (deflateInit2(&strm, aCompressionLevel, Z_DEFLATED, aWindowBits, aMemLevel, aStrategy) != Z_OK) {
        qWarning() << "Failed to initialize compression with custom settings.";
        return {};
    }

    if (!aDictionary.isEmpty()) {
        deflateSetDictionary(&strm, reinterpret_cast<const Bytef*>(aDictionary.constData()), aDictionary.size());
    }

    strm.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(aData.data()));
    strm.avail_in = aData.size();

    QByteArray compressed;
    char buffer[4096];

    int ret;
    do {
        strm.next_out = reinterpret_cast<Bytef*>(buffer);
        strm.avail_out = sizeof(buffer);

        ret = deflate(&strm, strm.avail_in ? Z_NO_FLUSH : Z_FINISH);
        if (ret != Z_OK && ret != Z_STREAM_END) {
            qWarning() << "Compression error:" << zError(ret);
            deflateEnd(&strm);
            return {};
        }

        compressed.append(buffer, sizeof(buffer) - strm.avail_out);
    } while (ret != Z_STREAM_END);

    deflateEnd(&strm);
    return compressed;
}

QByteArray Compression::DecompressDeflate(const QByteArray &aCompressedData) {
    if (aCompressedData.isEmpty())
        return {};

    z_stream strm{};
    strm.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(aCompressedData.data()));
    strm.avail_in = static_cast<uInt>(aCompressedData.size());

    // Negative window bits (-MAX_WBITS) indicate raw DEFLATE data.
    if (inflateInit2(&strm, -MAX_WBITS) != Z_OK) {
        qWarning() << "Failed to initialize DEFLATE for decompression.";
        return QByteArray();
    }

    QByteArray decompressed;
    char buffer[4096];

    int ret;
    do {
        strm.next_out = reinterpret_cast<Bytef*>(buffer);
        strm.avail_out = sizeof(buffer);

        ret = inflate(&strm, Z_NO_FLUSH);

        if (ret != Z_OK && ret != Z_STREAM_END) {
            qWarning() << "DEFLATE decompression error:" << zError(ret);
            inflateEnd(&strm);
            return QByteArray();
        }

        decompressed.append(buffer, sizeof(buffer) - strm.avail_out);
    } while (ret != Z_STREAM_END);

    inflateEnd(&strm);
    return decompressed;
}

QByteArray Compression::CompressDeflate(const QByteArray &aData) {
    return CompressDeflateWithSettings(aData);
}

QByteArray Compression::CompressDeflateWithSettings(const QByteArray &aData, int aCompressionLevel, int aWindowBits, int aMemLevel, int aStrategy, const QByteArray &aDictionary) {
    Q_UNUSED(aDictionary);

    if (aData.isEmpty())
        return QByteArray();

    z_stream strm{};

    // Negative window bits (-MAX_WBITS) indicate raw DEFLATE data.
    if (deflateInit2(&strm, aCompressionLevel, Z_DEFLATED, -aWindowBits, aMemLevel, aStrategy) != Z_OK) {
        qWarning() << "Failed to initialize DEFLATE for compression.";
        return QByteArray();
    }

    strm.next_in = reinterpret_cast<Bytef*>(const_cast<char*>(aData.data()));
    strm.avail_in = static_cast<uInt>(aData.size());

    QByteArray compressed;
    char buffer[4096];

    int ret;
    do {
        strm.next_out = reinterpret_cast<Bytef*>(buffer);
        strm.avail_out = sizeof(buffer);

        ret = deflate(&strm, strm.avail_in ? Z_NO_FLUSH : Z_FINISH);

        if (ret != Z_OK && ret != Z_STREAM_END) {
            qWarning() << "DEFLATE compression error:" << zError(ret);
            deflateEnd(&strm);
            return {};
        }

        compressed.append(buffer, sizeof(buffer) - strm.avail_out);
    } while (ret != Z_STREAM_END);

    deflateEnd(&strm);
    return compressed;
}

QByteArray Compression::DecompressLZO(const QByteArray &aCompressedData, quint32 aDestSize) {
    QByteArray dst;
    static bool ok = (lzo_init() == LZO_E_OK);
    if (!ok)
        throw std::runtime_error("lzo_init failed");

    dst = QByteArray(aDestSize, Qt::Uninitialized);
    lzo_uint out = aDestSize;

    int rc = lzo1x_decompress_safe(
        reinterpret_cast<const lzo_bytep>(aCompressedData.constData()),
        static_cast<lzo_uint>(aCompressedData.size()),
        reinterpret_cast<lzo_bytep>(dst.data()),
        &out,
        nullptr);

    if (rc != LZO_E_OK || out != aDestSize)
        throw std::runtime_error("LZO decompression error");

    return dst;
}

QByteArray Compression::DecompressOodle(const QByteArray &aCompressedData, quint32 aDecompressedSize) {
    return pDecompressOodle(aCompressedData, aCompressedData.length(), aDecompressedSize);
}

QByteArray Compression::CompressOodle(const QByteArray &aData) {
    quint32 maxSize = pGetOodleCompressedBounds(aData.length());
    QByteArray compressedData = pCompressOodle(aData, aData.length(),
                                               maxSize, OodleFormat::Kraken, OodleCompressionLevel::Optimal5);

    return compressedData.mid(0, maxSize);
}

quint32 Compression::pGetOodleCompressedBounds(quint32 aBufferSize) {
    return aBufferSize + 274 * ((aBufferSize + 0x3FFFF) / 0x400000);
}

QByteArray Compression::pCompressOodle(QByteArray aBuffer, quint32 aBufferSize, quint32 aOutputBufferSize, OodleFormat aformat, OodleCompressionLevel alevel) {
    QLibrary oodleLib("oo2core_8_win64");

    if (!oodleLib.load()) {
        qDebug() << "Failed to load DLL:" << oodleLib.errorString();
        return QByteArray();
    }

    OodleLZ_CompressFunc OodleLZ_Compress =
        (OodleLZ_CompressFunc)oodleLib.resolve("OodleLZ_Compress");

    if (!OodleLZ_Compress) {
        qDebug() << "Failed to resolve function:" << oodleLib.errorString();
        return QByteArray();
    }

    std::byte *outputBuffer = new std::byte[aOutputBufferSize];

    if (aBuffer.length() > 0 && aBufferSize > 0 && aOutputBufferSize > 0)
        OodleLZ_Compress(aformat, reinterpret_cast<std::byte*>(aBuffer.data()), aBufferSize, outputBuffer, alevel, 0, 0, 0);

    return QByteArray(reinterpret_cast<const char*>(outputBuffer), aOutputBufferSize);
}

QByteArray Compression::pDecompressOodle(QByteArray aBuffer, quint32 aBufferSize, quint32 aOutputBufferSize) {
    QLibrary oodleLib("oo2core_8_win64");

    if (!oodleLib.load()) {
        qDebug() << "Failed to load DLL:" << oodleLib.errorString();
        return QByteArray();
    }

    OodleLZ_DecompressFunc OodleLZ_Decompress =
        (OodleLZ_DecompressFunc)oodleLib.resolve("OodleLZ_Decompress");

    if (!OodleLZ_Decompress) {
        qDebug() << "Failed to resolve function:" << oodleLib.errorString();
        return QByteArray();
    }

    std::byte *outputBuffer = new std::byte[aOutputBufferSize];

    if (aBuffer.length() > 0 && aBufferSize > 0 && aOutputBufferSize > 0)
        OodleLZ_Decompress(reinterpret_cast<std::byte*>(aBuffer.data()), aBufferSize, outputBuffer, aOutputBufferSize, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);

    return QByteArray(reinterpret_cast<const char*>(outputBuffer), aOutputBufferSize);
}