/***
    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"

namespace Aurora::Compression
{
    static bool DecompressZSTD(const CompressionPipe &info)
    {
        auto length = ZSTD_DStreamInSize();

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

        auto outFrameLength = ZSTD_DStreamOutSize();
        AuList<AuUInt8> inflatedBuffer;
        ok = TryResize(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 = TryResize(inflatedBuffer, buffInSize);
        if (!ok)
        {
            SysPushErrorMem("Couldn't reserve deflation buffers");
            return false;
        }

        AuList<AuUInt8> deflatedBuffer;
        ok = TryResize(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, ArraySize(in));
            strm.avail_in = read;
            
            inputStat += read;

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

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

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

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

                stream.writePipe(out, have);

                outputStat += have;
                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, ArraySize(in));
            inputStat += read;
            if (!read)
            {
                break;
            }
            
            strm.avail_in = read;
            strm.next_in = in;

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

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

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

                stream.writePipe(out, have);

                outputStat += have;
                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, ArraySize(in));
            strm.avail_in = read;
            
            inputStat += read;

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

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

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

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

                stream.writePipe(out, have);

                outputStat += have;
                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, ArraySize(in));
            inputStat += read;
            if (!read)
            {
                break;
            }
            
            strm.avail_in = read;
            strm.next_in  = in;

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

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

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

                stream.writePipe(out, have);

                outputStat += have;
                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 = true;
        LZ4_stream_t* const lz4Stream = LZ4_createStream();
        char inBuf[kChunkSize];
        char outBuf[kChunkSize];
        AuUInt32 inputStat = 0, outputStat = 0;

        while (true)
        {
            auto read = stream.inPipe(inBuf, ArraySize(inBuf));
            if (!read) break;

            AuUInt16 bufferedBytes = LZ4_compress_fast_continue(lz4Stream, inBuf, outBuf, read, ArraySize(outBuf), 1);

            if (bufferedBytes <= 0)
            {
                ret = false;
                break;
            }
            
            stream.writePipe(&bufferedBytes, sizeof(bufferedBytes));
            stream.writePipe(outBuf, bufferedBytes);

            inputStat += read;
            outputStat += bufferedBytes;

            if (!stream.reportProgress(inputStat, outputStat)) break;
        }

        LZ4_freeStream(lz4Stream);
        return ret;
    }
    
    static  bool DecompressLZ4(const CompressionPipe &pipe)
    {
        bool ret = true;
        LZ4_streamDecode_t* lz4Stream = LZ4_createStreamDecode();
        char inBuf[kChunkSize];
        char outBuf[kChunkSize];
        AuUInt32 inputStat = 0, outputStat = 0;

        while (true)
        {
            AuUInt16 bufferedBytes;
            auto read = pipe.inPipe(&bufferedBytes, sizeof(bufferedBytes));
            if (read == 0)
            {
                break;
            }

            if (read != 2)
            {
                ret = false;
                goto out;
            }
            
            //while (bufferedBytes) TODO:
            {
                auto toRead = std::min(AuUInt16(ArraySize(inBuf)), bufferedBytes);
                
                read = pipe.inPipe(inBuf, toRead);
                if (read != toRead)
                {
                    ret = false;
                    goto out;
                }
                bufferedBytes -= toRead;
                
                auto err = LZ4_decompress_safe_continue(lz4Stream, inBuf, outBuf, read, ArraySize(outBuf));
                if (err < 0)
                {
                    ret = false;
                    goto out;
                }
                
                pipe.writePipe(outBuf, err);

                inputStat += bufferedBytes;
                outputStat += read;

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

        out:
        LZ4_freeStreamDecode(lz4Stream);
        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;
        }
    }
}