/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: StreamCompression.cpp
    Date: 2021-6-17
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Compression.hpp"
#include "StreamCompression.hpp"

#include "bzlib.h"
#include "zstd.h"
#include "zlib.h"
#include "lz4.h"
#include "lz4frame.h"

namespace Aurora::Compression
{
    static bool DecompressZSTD(const CompressionPipe &info)
    {
        Memory::ByteBuffer buffer;
        Memory::ByteBuffer inflatedBuffer;

        if (!info.writePipe)
        {
            return false;
        }

        if (!info.inPipe)
        {
            return false;
        }

        auto length = ZSTD_DStreamInSize();

        auto ok = AuTryResize(buffer, length);
        if (!ok)
        {
            SysPushErrorMem("Couldn't reserve inflation buffers");
            return false;
        }

        auto outFrameLength = ZSTD_DStreamOutSize();
        ok = AuTryResize(inflatedBuffer, outFrameLength);
        if (!ok)
        {
            SysPushErrorMem("Couldn't reserve inflation buffers");
            return false;
        }

        auto dctx = ZSTD_createDCtx();
        if (!dctx)
        {
            
            SysPushErrorGeneric("Couldn't create decompressor");
            return false;
        }

        AuUInt read = buffer.size();
        while (read = (info.inPipe(buffer.data(), length)))
        {
            ZSTD_inBuffer input = { buffer.data(), read, 0 };
            while (input.pos < input.size)
            {
                ZSTD_outBuffer output = { inflatedBuffer.data(), outFrameLength, 0 };
                auto ret = ZSTD_decompressStream(dctx, &output, &input);
                if (ZSTD_isError(ret))
                {
                    SysPushErrorIO("Compression error: {}", ret);
                    ZSTD_freeDCtx(dctx);
                    return false;
                }

                info.writePipe(output.dst, output.pos);
                if (info.reportProgress)
                {
                    if (!info.reportProgress(input.pos, output.pos))
                    {
                        ZSTD_freeDCtx(dctx);
                        return false;
                    }
                }
            }
        }

        ZSTD_freeDCtx(dctx);
        return true;
    }

    static bool CompressZSTD(const CompressionPipe &stream, const CompressionInfo &info)
    {
        Memory::ByteBuffer inflatedBuffer;
        Memory::ByteBuffer deflatedBuffer;
        size_t ret;
        const auto buffInSize = ZSTD_CStreamInSize();
        const auto buffOutSize = ZSTD_CStreamOutSize();
        
        if (!stream.writePipe)
        {
            return false;
        }

        if (!stream.inPipe)
        {
            return false;
        }

        auto ok = AuTryResize(inflatedBuffer, buffInSize);
        if (!ok)
        {
            SysPushErrorMem("Couldn't reserve deflation buffers");
            return false;
        }

        ok = AuTryResize(deflatedBuffer, buffOutSize);
        if (!ok)
        {
            SysPushErrorMem("Couldn't reserve deflation buffers. Out of memory");
            return false;
        }

        auto cctx = ZSTD_createCCtx();
        if (!cctx)
        {
            SysPushErrorGeneric("Couldn't create decompressor");
            return false;
        }

        ret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, info.compressionLevel);
        if (ZSTD_isError(ret))
        {
            SysPushErrorGen("Invalid compression level");
            ZSTD_freeCCtx(cctx);
            return false;
        }

        ret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_checksumFlag, 1);
        if (ZSTD_isError(ret))
        {
            SysPushErrorGen("Invalid option");
            ZSTD_freeCCtx(cctx);
            return false;
        }

        ret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, AuMax(stream.threads, 1u));
        if (ZSTD_isError(ret))
        {
            SysPushErrorGen();
            ZSTD_freeCCtx(cctx);
            return false;
        }

        bool status = true;

        size_t const toRead = buffInSize;

        while (true)
        {
            size_t read = stream.inPipe(inflatedBuffer.data(), buffInSize);
            /* Select the flush mode.
            * If the read may not be finished (read == toRead) we use
            * ZSTD_e_continue. If this is the last chunk, we use ZSTD_e_end.
            * Zstd optimizes the case where the first flush mode is ZSTD_e_end,
            * since it knows it is compressing the entire source in one pass.
             */
            int const lastChunk = (read < toRead);
            ZSTD_EndDirective const mode = lastChunk ? ZSTD_e_end : ZSTD_e_continue;
            /* Set the input buffer to what we just read.
            * We compress until the input buffer is empty, each time flushing the
            * output.
             */
            ZSTD_inBuffer input = { inflatedBuffer.data(), read, 0 };
            int finished;
            do
            {
                /* Compress into the output buffer and write all of the output to
                * the file so we can reuse the buffer next iteration.
                 */
                ZSTD_outBuffer output = { deflatedBuffer.data(), buffOutSize, 0 };
                size_t const remaining = ZSTD_compressStream2(cctx, &output, &input, mode);
                if (ZSTD_isError(remaining))
                {
                    SysPushErrorGen(": {}", ZSTD_getErrorName(remaining));
                    ZSTD_freeCCtx(cctx);
                    return false;
                }

                stream.writePipe(output.dst, output.pos);

                if (stream.reportProgress)
                {
                    if (!stream.reportProgress(input.pos, output.pos))
                    {
                        ZSTD_freeCCtx(cctx);
                        return false;
                    }
                }

                /* If we're on the last chunk we're finished when zstd returns 0,
                * which means its consumed all the input AND finished the frame.
                * Otherwise, we're finished when we've consumed all the input.
                 */
                finished = lastChunk ? (remaining == 0) : (input.pos == input.size);
            } while (!finished);

            SysAssertExp(input.pos == input.size);

            if (lastChunk) break;
        }

        ZSTD_freeCCtx(cctx);
        return true;
    }

    static bool CompressZLib(const CompressionPipe &stream, const CompressionInfo &info)
    {
        int ret, flush;
        z_stream strm {};
        unsigned char in[kChunkSize];
        unsigned char out[kChunkSize];
        
        if (!stream.writePipe)
        {
            return false;
        }

        if (!stream.inPipe)
        {
            return false;
        }

        ret = deflateInit(&strm, info.compressionLevel);
        if (ret < Z_OK)
        {
            SysPushErrorIO("Error: {}", ret);
            return false;
        }
        
        AuUInt32 inputStat = 0, outputStat = 0;

        do 
        {
            auto read = stream.inPipe(in, AuArraySize(in));
            strm.avail_in = read;
            
            inputStat += read;

            flush = !strm.avail_in ? Z_FINISH : Z_NO_FLUSH;
            strm.next_in = in;

            do
            {
                strm.avail_out = AuArraySize(out);
                strm.next_out  = out;

                ret = deflate(&strm, flush);
                if (ret < Z_OK)
                {
                    SysPushErrorIO("Error: {}", ret);
                    return false;
                }

                auto have = AuArraySize(out) - strm.avail_out;

                stream.writePipe(out, have);

                outputStat += have;
                if (stream.reportProgress)
                {
                    if (!stream.reportProgress(inputStat, outputStat))
                    {
                        deflateEnd(&strm);
                        return false;
                    }
                }

            } while (strm.avail_out == 0);
            SysAssert(strm.avail_in == 0);

        } while (flush != Z_FINISH);

        deflateEnd(&strm);
        return true;
    }
    
    static bool DecompressZLib(const CompressionPipe &stream)
    {
        int ret;
        z_stream strm {};
        unsigned char in[kChunkSize];
        unsigned char out[kChunkSize];
        
        if (!stream.writePipe)
        {
            return false;
        }

        if (!stream.inPipe)
        {
            return false;
        }

        ret = inflateInit(&strm);
        if (ret < Z_OK)
        {
            SysPushErrorIO("Error: {}", ret);
            return false;
        }
        
        AuUInt32 inputStat = 0, outputStat = 0;

        do 
        {
            auto read = stream.inPipe(in, AuArraySize(in));
            inputStat += read;
            if (!read)
            {
                break;
            }
            
            strm.avail_in = read;
            strm.next_in = in;

            do
            {
                strm.avail_out = AuArraySize(out);
                strm.next_out  = out;

                ret = inflate(&strm, Z_NO_FLUSH);
                if (ret < Z_OK)
                {
                    SysPushErrorIO("Error: {}", ret);
                    inflateEnd(&strm);
                    return false;
                }

                auto have = AuArraySize(out) - strm.avail_out;

                stream.writePipe(out, have);

                outputStat += have;
                if (stream.reportProgress)
                {
                    if (!stream.reportProgress(inputStat, outputStat))
                    {
                        inflateEnd(&strm);
                        return false;
                    }
                }

            } while (strm.avail_out == 0);
            SysAssert(strm.avail_in == 0);

        } while (ret != Z_STREAM_END);

        inflateEnd(&strm);
        return true;
    }

    static bool CompressBZip2(const CompressionPipe &stream, const CompressionInfo &info)
    {
        int ret, flush;
        bz_stream strm {};
        char in[kChunkSize];
        char out[kChunkSize];
        
        if (!stream.writePipe)
        {
            return false;
        }

        if (!stream.inPipe)
        {
            return false;
        }

        ret = BZ2_bzCompressInit(&strm, info.compressionLevel, 0, 0);
        if (ret < BZ_OK)
        {
            SysPushErrorIO("Error: {}", ret);
            return false;
        }
        
        AuUInt32 inputStat = 0, outputStat = 0;

        do 
        {
            auto read = stream.inPipe(in, AuArraySize(in));
            strm.avail_in = read;
            
            inputStat += read;

            flush = !strm.avail_in ? BZ_FINISH : BZ_RUN;
            strm.next_in = in;

            do
            {
                strm.avail_out = AuArraySize(out);
                strm.next_out  = out;

                ret = BZ2_bzCompress(&strm, flush);
                if (ret < BZ_OK)
                {
                    SysPushErrorIO("Error: {}", ret);
                    BZ2_bzCompressEnd(&strm);
                    return false;
                }

                auto have = AuArraySize(out) - strm.avail_out;

                stream.writePipe(out, have);

                outputStat += have;

                if (stream.reportProgress)
                {
                    if (!stream.reportProgress(inputStat, outputStat))
                    {
                        BZ2_bzCompressEnd(&strm);
                        return false;
                    }
                }

            } while (strm.avail_out == 0);
            SysAssert(strm.avail_in == 0);

        } while (flush != BZ_FINISH);

        BZ2_bzCompressEnd(&strm);
        return true;
    }

    static bool DecompressBZip2(const CompressionPipe &stream)
    {
        int ret;
        bz_stream strm {};
        char in[kChunkSize];
        char out[kChunkSize];
        
        if (!stream.writePipe)
        {
            return false;
        }

        if (!stream.inPipe)
        {
            return false;
        }

        ret = BZ2_bzDecompressInit(&strm, 0, 0);
        if (ret < Z_OK)
        {
            SysPushErrorIO("Error: {}", ret);
            return false;
        }
        
        AuUInt32 inputStat = 0, outputStat = 0;

        do 
        {
            auto read = stream.inPipe(in, AuArraySize(in));
            inputStat += read;
            if (!read)
            {
                break;
            }
            
            strm.avail_in = read;
            strm.next_in  = in;

            do
            {
                strm.avail_out = AuArraySize(out);
                strm.next_out  = out;

                ret = BZ2_bzDecompress(&strm);
                if (ret < Z_OK)
                {
                    SysPushErrorIO("Error: {}", ret);
                    BZ2_bzDecompressEnd(&strm);
                    return false;
                }

                auto have = AuArraySize(out) - strm.avail_out;

                stream.writePipe(out, have);

                outputStat += have;
                if (stream.reportProgress)
                {
                    if (!stream.reportProgress(inputStat, outputStat))
                    {
                        BZ2_bzDecompressEnd(&strm);
                        return false;
                    }
                }

            } while (strm.avail_out == 0);
            SysAssert(strm.avail_in == 0);

        } while (ret != BZ_FINISH);

        BZ2_bzDecompressEnd(&strm);
        return true;
    }

    static bool CompressLZ4(const CompressionPipe &stream, const CompressionInfo &info)
    {
        bool ret;
        LZ4F_cctx *cctxPtr;
        LZ4F_preferences_t pref {};
        AuUInt32 inputStat = 0, outputStat = 0;
        LZ4F_compressOptions_t options {};
        char header[512];
        
        ret = true;
        pref.compressionLevel = info.compressionLevel;

        auto maxFrameSize = info.lz4BlockSize ? info.lz4BlockSize : 64 * 1024;
        auto buffer       = AuSPtr<char>(new char[maxFrameSize], AuDefaultDeleter<char[]>());
        auto maxOut       = LZ4F_compressBound(maxFrameSize, &pref);
        auto outBuffer    = AuSPtr<char>(new char[maxOut], AuDefaultDeleter<char[]>());;

        auto err = LZ4F_createCompressionContext(&cctxPtr, LZ4F_getVersion());
        if (LZ4F_isError(err))
        {
            return false;
        }

        // Write header
        {
            auto written = LZ4F_compressBegin(cctxPtr, header, AuArraySize(header), &pref);
            stream.writePipe(header, written);
        }

        while (true)
        {
            auto read = stream.inPipe(buffer.get(), maxFrameSize);
            inputStat += read;

            if (read)
            {
                AuUInt32 bufferedBytes = LZ4F_compressUpdate(cctxPtr, outBuffer.get(), maxOut, buffer.get(), read, &options);

                if (LZ4F_isError(bufferedBytes))
                {
                    ret = false;
                    break;
                }

                stream.writePipe(outBuffer.get(), bufferedBytes);
                
                outputStat += bufferedBytes;
            }
            else
            {
                AuUInt32 bufferedBytes = LZ4F_compressEnd(cctxPtr, outBuffer.get(), maxOut, &options);
                stream.writePipe(outBuffer.get(), bufferedBytes);
                break;
            }

            if (stream.reportProgress)
            {
                if (!stream.reportProgress(inputStat, outputStat))
                {
                    ret = false;
                    break;
                }
            }
        }

        LZ4F_freeCompressionContext(cctxPtr);
        return ret;
    }

    static  bool DecompressLZ4(const CompressionPipe &pipe, AuUInt32 frameMaxSize)
    {
        bool ret = true;
        LZ4F_dctx *dctxPtr;
        AuUInt32 inputStat = 0, outputStat = 0;
        size_t bytesRemInFrame {};
        LZ4F_decompressOptions_t opts {};
        AuSPtr<char> bufferIn, bufferOut;
        
        if (!pipe.writePipe)
        {
            return false;
        }

        if (!pipe.inPipe)
        {
            return false;
        }

        auto err = LZ4F_createDecompressionContext(&dctxPtr, LZ4F_getVersion());
        if (LZ4F_isError(err))
        {
            return false;
        }
        bufferIn  = AuSPtr<char>(new char[frameMaxSize / 2], AuDefaultDeleter<char[]>());
        bufferOut = AuSPtr<char>(new char[frameMaxSize / 2], AuDefaultDeleter<char[]>());

        while (true)
        {
            auto frameSize = bytesRemInFrame ? bytesRemInFrame : LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH;
            auto min = frameSize;

            if (min > (frameMaxSize / 2))
            {
                ret = false;
                break;
            }

            if (frameSize)
            {
                auto read = pipe.inPipe(bufferIn.get(), frameSize);
                if (read != frameSize)
                {
                    ret = read == 0 && outputStat;
                    break;
                }
                inputStat += frameSize;
            }

            if (frameSize)
            {
                auto mustConsume = frameSize;
                size_t frameSPtr = mustConsume;
                size_t frameS2Ptr = frameMaxSize;

                bytesRemInFrame = LZ4F_decompress(dctxPtr, bufferOut.get(), &frameS2Ptr, bufferIn.get(), &frameSPtr, &opts);
                if (LZ4F_isError(bytesRemInFrame))
                {
                    ret = false;
                    break;
                }
                
                if (frameS2Ptr)
                {
                    pipe.writePipe(bufferOut.get(), frameS2Ptr);
                }

                outputStat += frameS2Ptr;

                if (pipe.reportProgress)
                {
                    if (!pipe.reportProgress(inputStat, outputStat))
                    {
                        ret = false;
                        break;
                    }
                }
            }
        }

        LZ4F_freeDecompressionContext(dctxPtr);
        return ret;
    }

    AUKN_SYM bool Compress(const CompressionPipe &stream, const CompressionInfo &info)
    {
        if (!stream.inPipe)
        {
            return {};
        }

        if (!stream.writePipe)
        {
            return {};
        }

        switch (info.type)
        {
        case ECompresionType::eZSTD:
            return CompressZSTD(stream, info);
        case ECompresionType::eDeflate:
            return CompressZLib(stream, info);
        case ECompresionType::eBZIP2:
            return CompressBZip2(stream, info);
        case ECompresionType::eLZ4:
            return CompressLZ4(stream, info);
        //case ECompresionType::eLZMA:
        //    return CompressLZMA(stream, info);
        default:
            return false;
        }
    }

    AUKN_SYM bool Decompress(const CompressionPipe &pipe, const DecompressInfo &meta2)
    {  
        DecompressInfo meta = meta2;
        
        if (!meta.internalStreamSize)
        {
            meta.internalStreamSize = 1024 * 64 * 2;
        }

        if (!pipe.inPipe)
        {
            return {};
        }

        if (!pipe.writePipe)
        {
            return {};
        }

        switch (meta.alg)
        {
        case ECompresionType::eZSTD:
            return DecompressZSTD(pipe);
        case ECompresionType::eDeflate:
            return DecompressZLib(pipe);
        case ECompresionType::eBZIP2:
            return DecompressBZip2(pipe);
        case ECompresionType::eLZ4:
            return DecompressLZ4(pipe, meta.internalStreamSize);
        //case ECompresionType::eLZMA:
        //    return DecompressLZMA(pipe);
        default:
            return false;
        }
    }
}