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

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

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

namespace Aurora::Compression
{
    bool BaseStream::Read(void * /*opt*/ buffer, AuUInt32 &len, bool ingestUntilError)
    {
        if (ingestUntilError)
        {
            while (this->_outbuffer.size() < len)
            {
                if (!Ingest(0, len))
                {
                    return false;
                }
            }
        }

        return StreamRead(buffer, len, this->_outbuffer);
    }

    class ZSTDInflate : public BaseStream
    {
    public:
        ~ZSTDInflate()
        {
            if (auto dctx = std::exchange(_dctx, {}))
            {
                ZSTD_freeDCtx(dctx);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->_reader = reader;
            this->_dctx = ZSTD_createDCtx();
            
            if (!this->_dctx)
            {
                SysPushErrorGen("Couldn't create decompressor");
                return false;
            }

            this->_outbuffer.reserve(10 * 1024);

            return true;
        }

        bool Ingest(AuUInt32 minimum, AuUInt32 requestBuffer) override
        {
            AuUInt32 length = ZSTD_DStreamInSize();
            void *din = alloca(length);

            AuUInt32 readLength = requestBuffer;
            AuUInt32 requested = std::min(readLength, length);
            AuUInt32 request = requested;
            AuUInt32 done{};

            auto outFrameLength = ZSTD_DStreamOutSize();
            void *dout = alloca(outFrameLength);

            while (done != readLength)
            {
                if (this->_reader->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return minimum <= done;
                }

                done += request;

                ZSTD_inBuffer input = { din, request, 0 };
                while (input.pos < input.size)
                {
                    ZSTD_outBuffer output = { dout, outFrameLength, 0 };
                    
                    auto ret = ZSTD_decompressStream(this->_dctx, &output, &input);
                    if (ZSTD_isError(ret))
                    {
                        SysPushErrorIO("Compression error: {}", ret);
                        return false;
                    }

                    this->_outbuffer.insert(this->_outbuffer.end(), reinterpret_cast<const AuUInt8 *>(output.dst), reinterpret_cast<const AuUInt8 *>(output.dst) + output.pos);
                }
            }

            return true;
        }
  
    private:
        Aurora::IO::IStreamReader *_reader;
        ZSTD_DCtx *_dctx;

    };

    class ZIPInflate : public BaseStream
    {
    public:
        ~ZIPInflate()
        {
            if (auto ctx = std::exchange(this->_init, {}))
            {
                inflateEnd(&this->_ctx);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->_reader = reader;
            
            auto ret = inflateInit(&this->_ctx);
            if (ret != Z_OK)
            {
                SysPushErrorMem("Error: {}", ret);
                return false;
            }

            this->_outbuffer.reserve(10 * 1024);
            this->_init = true;
            return true;
        }

        bool Ingest(AuUInt32 minimum, AuUInt32 requestBuffer) override
        {
            int ret;
            AuUInt32 readLength = requestBuffer;
            unsigned char dout[4096];
            unsigned char din[4096];
            AuUInt32 request = std::min(requestBuffer, AuUInt32(ArraySize(din)));
            
            AuUInt32 done{};
            while (done != readLength)
            {
                if (this->_reader->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return minimum <= done;
                }

                done += request;

                this->_ctx.avail_in = request;
                this->_ctx.next_in  = reinterpret_cast<unsigned char *>(din);

                do
                {
                    this->_ctx.avail_out = ArraySize(dout);
                    this->_ctx.next_out  = dout;

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

                    auto have = ArraySize(dout) - this->_ctx.avail_out;

                    this->_outbuffer.insert(this->_outbuffer.end(), reinterpret_cast<const AuUInt8 *>(dout), reinterpret_cast<const AuUInt8 *>(dout) + have);


                } while (this->_ctx.avail_out == 0);
                SysAssert(this->_ctx.avail_in == 0);
            }

            return true;
        }

    private:
        AuList<AuUInt8> _outbuffer;
        Aurora::IO::IStreamReader *_reader;
        z_stream _ctx {};
        bool _init {};

    };

    class BZIPInflate : public BaseStream
    {
    public:
        ~BZIPInflate()
        {
            if (auto ctx = std::exchange(this->_init, {}))
            {
                BZ2_bzDecompressEnd(&this->_ctx);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->_reader = reader;
            
            auto ret = BZ2_bzDecompressInit(&this->_ctx, 0, 0);
            if (ret != Z_OK)
            {
                SysPushErrorMem("Error: {}", ret);
                return false;
            }

            this->_outbuffer.reserve(10 * 1024);
            this->_init = true;
            return true;
        }

        bool Ingest(AuUInt32 minimum, AuUInt32 requestBuffer) override
        {
            int ret;
            char dout[4096];
            char din[4096];

            AuUInt32 readLength = requestBuffer;
            AuUInt32 request = std::min(requestBuffer, AuUInt32(ArraySize(din)));
            
            AuUInt32 done{};
            while (done != readLength)
            {
                if (this->_reader->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return minimum <= done;
                }

                done += request;

                this->_ctx.avail_in = request;
                this->_ctx.next_in  = reinterpret_cast<char *>(din);

                do
                {
                    this->_ctx.avail_out = ArraySize(dout);
                    this->_ctx.next_out  = dout;

                    ret = BZ2_bzDecompress(&this->_ctx);
                    if (ret != Z_OK)
                    {
                        SysPushErrorIO("Error: {}", ret);
                        return false;
                    }

                    auto have = ArraySize(dout) - this->_ctx.avail_out;

                    this->_outbuffer.insert(this->_outbuffer.end(), reinterpret_cast<const AuUInt8 *>(dout), reinterpret_cast<const AuUInt8 *>(dout) + have);


                } while (this->_ctx.avail_out == 0);
                SysAssert(this->_ctx.avail_in == 0);
            }

            return true;
        }

    private:
        AuList<AuUInt8> _outbuffer;
        Aurora::IO::IStreamReader *_reader;
        bz_stream _ctx {};
        bool _init {};

    };

    class LZ4Inflate : public BaseStream
    {
    public:
        ~LZ4Inflate()
        {
            if (auto ctx = std::exchange(this->_lz4Stream, {}))
            {
                LZ4_freeStreamDecode(this->_lz4Stream);
            }
        }

        bool Init(Aurora::IO::IStreamReader *reader)
        {
            this->_reader = reader;
            
            this->_lz4Stream = LZ4_createStreamDecode();
            if (!this->_lz4Stream)
            {
                SysPushErrorMem();
                return false;
            }

            this->_outbuffer.reserve(10 * 1024);
            return true;
        }

        bool Ingest(AuUInt32 minimum, AuUInt32 requestBuffer) override
        {
            int ret;
            char dout[4096];
            char din[4096];
            
            AuUInt32 readLength = requestBuffer;
            AuUInt32 request = std::min(requestBuffer, AuUInt32(ArraySize(din)));
            AuUInt32 done{};

            while (done != readLength)
            {
                if (this->_reader->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return minimum <= done;
                }

                done += request;

             
                int framesLength = done;
                char *framesPointer = (char *)din;
                while (true)
                {
                    if (framesLength < 2)
                    {
                        return false;
                    }
                    framesLength -= 2;

                    auto frameLength = *reinterpret_cast<AuUInt16 *>(framesPointer);
                    if (frameLength > framesLength)
                    {
                        return false;
                    }
                    framesPointer += 2;

                    auto err = LZ4_decompress_safe_continue(_lz4Stream, framesPointer, dout, frameLength, ArraySize(dout));
                    if (err < 0)
                    {
                        ret = false;
                        return false;
                    }
                    
                    this->_outbuffer.insert(this->_outbuffer.end(), reinterpret_cast<const AuUInt8 *>(dout), reinterpret_cast<const AuUInt8 *>(dout) + err);

                    framesPointer += framesLength;
                    framesLength  -= framesLength;
                }
            }

            return true;
        }

    private:
        AuList<AuUInt8> _outbuffer;
        Aurora::IO::IStreamReader *_reader;
        LZ4_streamDecode_t* _lz4Stream {};

    };

    AUKN_SYM ICompressionStream *DecompressorNew(IO::IStreamReader *reader, ECompresionType type)
    {
        BaseStream * ret{};

        switch (type)
        {
        case ECompresionType::eZSTD:
            ret = new ZSTDInflate();
            break;
        case ECompresionType::eBZIP2:
            ret = new BZIPInflate();
            break;
        case ECompresionType::eLZ4:
            ret = new LZ4Inflate();
            break;
        case ECompresionType::eDeflate:
            ret = new ZIPInflate();
            break;
        default:
            ret = nullptr;
            break;
        }
        
        if (ret)
        {
            if (!ret->Init(reader))
            {
                ret = nullptr;
                delete ret;
            }
        }

        return ret;
    }

    AUKN_SYM void DecompressorRelease(ICompressionStream * stream)
    {
        SafeDelete<BaseStream *>(stream);
    }
}