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

    File: StreamCompression.cpp
    Date: 2021-6-17
    Author: Reece
***/
#include <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)
    {
        auto length = ZSTD_DStreamInSize();

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

        auto outFrameLength = ZSTD_DStreamOutSize();
        AuList<AuUInt8> inflatedBuffer;
        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)
    {
        size_t ret;
        const auto buffInSize = ZSTD_CStreamInSize();
        const auto buffOutSize = ZSTD_CStreamOutSize();

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

        AuList<AuUInt8> deflatedBuffer;
        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, std::max(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];

        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];

        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);
                    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];

        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);
                    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];

        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);
                    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], std::default_delete<char[]>());
        auto maxOut       = LZ4F_compressBound(maxFrameSize, &pref);
        auto outBuffer    = AuSPtr<char>(new char[maxOut], std::default_delete<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)
    {
        bool ret = true;
        LZ4F_dctx *dctxPtr;
        AuUInt32 inputStat = 0, outputStat = 0;
        LZ4F_decompressOptions_t opts {};
        AuUInt32 lastFrameMaxSize {};
        AuSPtr<char> bufferIn, bufferOut;

        auto err = LZ4F_createDecompressionContext(&dctxPtr, LZ4F_getVersion());
        if (LZ4F_isError(err))
        {
            return false;
        }

        while (true)
        {
            char header[LZ4F_HEADER_SIZE_MAX];

            LZ4F_frameInfo_t info {};
             
            // Read header
            if (pipe.inPipe(header, LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH) != LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH)
            {
                return {};
            }

            // 
            auto sizeOfheader = LZ4F_headerSize(header, LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH);
            if (LZ4F_isError(sizeOfheader))
            {
                return {};
            }

            const auto req = sizeOfheader - LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH;
            if (pipe.inPipe(header + LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH, req) != req)
            {
                return {};
            }
            inputStat += sizeOfheader;

            size_t sizePtr = sizeOfheader;
            auto status = LZ4F_getFrameInfo(dctxPtr, &info, header, &sizePtr);
            if (sizePtr != sizeOfheader)
            {
                return {};
            }

            auto frameSize = status;
            auto min = std::max(info.contentSize, frameSize);

            if (min > lastFrameMaxSize)
            {
                lastFrameMaxSize = min;
                if (lastFrameMaxSize > 64 * 1024 * 1024)
                {
                    return {};
                }
                bufferIn  = AuSPtr<char>(new char[lastFrameMaxSize], std::default_delete<char[]>());
                bufferOut = AuSPtr<char>(new char[lastFrameMaxSize], std::default_delete<char[]>());
            }

            if (pipe.inPipe(bufferIn.get(), frameSize) != frameSize)
            {
                return {};
            }

            inputStat += frameSize;

            if (info.contentSize)
            {
                size_t frameSPtr = frameSize;
                size_t frameS2Ptr = info.contentSize;
                LZ4F_decompress(dctxPtr, bufferIn.get(), &frameSPtr, bufferOut.get(), &frameS2Ptr, &opts);

                if (frameS2Ptr)
                {
                    pipe.writePipe(bufferOut.get(), frameS2Ptr);
                }

                outputStat += frameS2Ptr;

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

        LZ4F_freeDecompressionContext(dctxPtr);
        return ret;
    }

    AUKN_SYM bool Compress(const CompressionPipe &stream, const CompressionInfo &info)
    {
        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)
    {  
        switch (pipe.type)
        {
        case ECompresionType::eZSTD:
            return DecompressZSTD(pipe);
        case ECompresionType::eDeflate:
            return DecompressZLib(pipe);
        case ECompresionType::eBZIP2:
            return DecompressBZip2(pipe);
        case ECompresionType::eLZ4:
            return DecompressLZ4(pipe);
        //case ECompresionType::eLZMA:
        //    return DecompressLZMA(pipe);
        default:
            return false;
        }
    }
}