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

#include <QLibrary>
#include <QDebug>
#include <QFile>
#include <QDir>
#include <QProcess>
#include <QtEndian>

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.size() < 8)
        return {};

    // Read header identifier (big-endian in file)
    quint32 identifier = qFromBigEndian<quint32>(reinterpret_cast<const uchar*>(data.constData()));

    if (identifier == 0x0FF512EE) {
        // LZXNATIVE format - full header with codec params and block structure
        return DecompressXMemNative(data);
    } else if (identifier == 0x0FF512ED) {
        // LZXTDECODE format - simpler format
        return DecompressXMemTDecode(data, flags, windowSize, partSize);
    } else {
        // No header - treat as raw LZX data with provided params
        return DecompressXMemRaw(data, flags, windowSize, partSize);
    }
}

QByteArray Compression::DecompressXMemNative(const QByteArray &data)
{
    if (data.size() < 40) {
        qWarning() << "LZXNATIVE header too short";
        return {};
    }

    const uchar* ptr = reinterpret_cast<const uchar*>(data.constData());

    // Parse XCOMPRESS_FILE_HEADER_LZXNATIVE (all fields big-endian in file)
    // Offset 0-3: Identifier (already validated)
    // Offset 4-5: Version
    // Offset 6-7: Reserved
    // Offset 8-11: ContextFlags
    // Offset 12-15: CodecParams.Flags
    // Offset 16-19: CodecParams.WindowSize
    // Offset 20-23: CodecParams.CompressionPartitionSize
    // Offset 24-27: UncompressedSizeHigh
    // Offset 28-31: UncompressedSizeLow
    // Offset 32-35: CompressedSizeHigh
    // Offset 36-39: CompressedSizeLow
    // Offset 40-43: UncompressedBlockSize
    // Offset 44-47: CompressedBlockSizeMax

    XMEMCODEC_PARAMETERS_LZX lzxParams = {};
    lzxParams.Flags = qFromBigEndian<quint32>(ptr + 12);
    lzxParams.WindowSize = qFromBigEndian<quint32>(ptr + 16);
    lzxParams.CompressionPartitionSize = qFromBigEndian<quint32>(ptr + 20);

    quint64 uncompressedSize = (static_cast<quint64>(qFromBigEndian<quint32>(ptr + 24)) << 32) |
                                qFromBigEndian<quint32>(ptr + 28);
    quint32 uncompressedBlockSize = qFromBigEndian<quint32>(ptr + 40);

    qDebug() << "LZXNATIVE: windowSize=" << lzxParams.WindowSize
             << "partSize=" << lzxParams.CompressionPartitionSize
             << "uncompressedSize=" << uncompressedSize
             << "blockSize=" << uncompressedBlockSize;

    // Create decompression context
    XMEMDECOMPRESSION_CONTEXT ctx = nullptr;
    HRESULT hr = XMemCreateDecompressionContext(XMEMCODEC_LZX, &lzxParams, 0, &ctx);
    if (FAILED(hr) || !ctx) {
        qWarning() << "Failed to create LZX decompression context, hr=" << Qt::hex << hr;
        return {};
    }

    QByteArray output;
    output.reserve(static_cast<int>(uncompressedSize));

    // Data starts after 48-byte header, in blocks prefixed by 4-byte size
    int offset = 48;
    while (offset + 4 <= data.size()) {
        quint32 compressedBlockSize = qFromBigEndian<quint32>(ptr + offset);
        offset += 4;

        if (compressedBlockSize == 0 || offset + static_cast<int>(compressedBlockSize) > data.size())
            break;

        // Decompress this block
        QByteArray blockOut(uncompressedBlockSize, Qt::Uninitialized);
        SIZE_T destSize = blockOut.size();
        SIZE_T srcSize = compressedBlockSize;

        hr = XMemDecompress(ctx, blockOut.data(), &destSize, ptr + offset, srcSize);
        if (FAILED(hr)) {
            qWarning() << "XMemDecompress block failed, hr=" << Qt::hex << hr;
            XMemDestroyDecompressionContext(ctx);
            return output; // Return what we have
        }

        output.append(blockOut.constData(), static_cast<int>(destSize));
        offset += compressedBlockSize;

        // Reset context for next block
        XMemResetDecompressionContext(ctx);
    }

    XMemDestroyDecompressionContext(ctx);
    return output;
}

QByteArray Compression::DecompressXMemTDecode(const QByteArray &data, int flags, int windowSize, int partSize)
{
    Q_UNUSED(flags);
    Q_UNUSED(windowSize);
    Q_UNUSED(partSize);

    if (data.size() < 8)
        return {};

    // LZXTDECODE (0x0FF512ED) - Use QuickBMS as external tool since xcompress64.dll is unreliable
    // Write temp file, run QuickBMS, read result

    QString tempDir = QDir::tempPath();
    QString inputFile = tempDir + "/xmem_input.bin";
    QString outputFile = tempDir + "/xmem_output.bin";
    QString bmsScript = tempDir + "/xmem_decomp.bms";

    // Write input data
    QFile inFile(inputFile);
    if (!inFile.open(QIODevice::WriteOnly)) {
        qWarning() << "Failed to create temp input file";
        return {};
    }
    inFile.write(data);
    inFile.close();

    // Write BMS script
    QFile scriptFile(bmsScript);
    if (!scriptFile.open(QIODevice::WriteOnly | QIODevice::Text)) {
        qWarning() << "Failed to create BMS script";
        return {};
    }
    scriptFile.write("comtype xmemdecompress\n"
                     "get SIZE asize\n"
                     "clog \"\" 0 SIZE SIZE\n");
    scriptFile.close();

    // Run QuickBMS
    QProcess proc;
    proc.setWorkingDirectory(tempDir);
    QString quickbmsPath = "E:/Software/QuickBMS/quickbms.exe";
    QStringList args = {"-o", "-O", outputFile, bmsScript, inputFile};

    proc.start(quickbmsPath, args);
    if (!proc.waitForFinished(30000)) {
        qWarning() << "QuickBMS timeout or failed to start";
        return {};
    }

    if (proc.exitCode() != 0) {
        qWarning() << "QuickBMS failed:" << proc.readAllStandardError();
        return {};
    }

    // Read output
    QFile outFile(outputFile);
    if (!outFile.open(QIODevice::ReadOnly)) {
        qWarning() << "Failed to read QuickBMS output";
        return {};
    }
    QByteArray result = outFile.readAll();
    outFile.close();

    // Cleanup
    QFile::remove(inputFile);
    QFile::remove(outputFile);
    QFile::remove(bmsScript);

    qDebug() << "TDECODE via QuickBMS:" << data.size() << "->" << result.size() << "bytes";
    return result;
}

QByteArray Compression::DecompressXMemRaw(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;

    XMEMDECOMPRESSION_CONTEXT ctx = nullptr;
    HRESULT hr = XMemCreateDecompressionContext(XMEMCODEC_LZX, &lzxParams, 0, &ctx);
    if (FAILED(hr) || !ctx) {
        qWarning() << "Failed to create raw LZX context";
        return {};
    }

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

    QByteArray output;
    output.reserve(16 * 1024 * 1024);

    QByteArray scratch(0x10000, Qt::Uninitialized);

    while (availIn > 0) {
        SIZE_T inSize = availIn;
        SIZE_T outSize = scratch.size();

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

        if (FAILED(hr)) {
            qWarning() << "XMemDecompressStream raw failed, hr=" << Qt::hex << hr;
            break;
        }

        if (inSize == 0 && outSize == 0)
            break;

        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.size(), 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 lib("../../../third_party/oodle_lib/dll/oo2core_8_win64.dll"); // adjust path if needed
    if (!lib.load()) {
        qDebug() << "Failed to load:" << lib.errorString();
        return QByteArray();
    }

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

    if (!OodleLZ_Compress) {
        qDebug() << "Failed to resolve OodleLZ_Compress:" << lib.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(const QByteArray &aBuffer,
                                         quint32 aBufferSize,
                                         quint32 aOutputBufferSize)
{
    QLibrary lib("../../../third_party/oodle_lib/dll/oo2core_8_win64.dll");
    if (!lib.load()) {
        qWarning() << "Failed to load Oodle DLL:" << lib.errorString();
        return {};
    }

    OodleLZ_DecompressFunc OodleLZ_Decompress =
        reinterpret_cast<OodleLZ_DecompressFunc>(lib.resolve("OodleLZ_Decompress"));
    if (!OodleLZ_Decompress) {
        qWarning() << "Failed to resolve OodleLZ_Decompress:" << lib.errorString();
        return {};
    }

    QByteArray out(aOutputBufferSize + 1, Qt::Uninitialized);

    if (aBuffer.isEmpty() || aBufferSize == 0 || aOutputBufferSize == 0) {
        qWarning() << "Invalid Oodle parameters (empty input or size 0)";
        return {};
    }

    int result = OodleLZ_Decompress(
        aBuffer.constData(),
        static_cast<int64_t>(aBufferSize),
        out.data(),
        static_cast<int64_t>(aOutputBufferSize),
        1,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        3);

    if (result < 0) {
        qWarning() << "OodleLZ_Decompress failed with code" << result;
        return {};
    }

    if (result > out.size()) {
        qWarning() << "Oodle returned more than expected:" << result
                    << "expected" << aOutputBufferSize;
        return out;
    }

    out.resize(result);

    return out;
}