/***
    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(4096).second == 0)
                {
                    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;
        }

        std::pair<AuUInt32, AuUInt32> Ingest(AuUInt32 input) override
        {
            AuUInt32 length = ZSTD_DStreamInSize();
            void *din = alloca(length);
            auto outFrameLength = ZSTD_DStreamOutSize();
            void *dout = alloca(outFrameLength);
            AuUInt32 done{}, read{};

            while (read != input)
            {
                AuUInt32 request = std::min(input, length);
                if (this->_reader->Read(din, request) != IO::EStreamError::eErrorNone)
                {
                    return std::make_pair(read, done);
                }
                read += 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 std::make_pair(read, 0);
                    }

                    done += output.pos;
                    this->_outbuffer.insert(this->_outbuffer.end(),
                                            reinterpret_cast<const AuUInt8 *>(output.dst),
                                            reinterpret_cast<const AuUInt8 *>(output.dst) + output.pos);
                }
            }
            
            return std::make_pair(read, done);
        }
  
    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;
        }

        std::pair<AuUInt32, AuUInt32> Ingest(AuUInt32 input) override
        {
            int ret;

            AuUInt32 done{}, read{};
            
            while (read != input)
            {
                AuUInt32 request = std::min(input, AuUInt32(ArraySize(din_)));
                if (this->_reader->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return std::make_pair(read, done);
                }

                read += 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_;

                    if (!this->_ctx.avail_out)
                    {
                        break;
                    }

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

                    auto have = ArraySize(dout_) - this->_ctx.avail_out;
                    done += have;

                    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 std::make_pair(read, done);
        }

    private:
        AuList<AuUInt8> _outbuffer;
        Aurora::IO::IStreamReader *_reader;
        z_stream _ctx {};
        bool _init {};
        unsigned char din_[4096];
        unsigned char dout_[4096];

    };

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

        std::pair<AuUInt32, AuUInt32> Ingest(AuUInt32 input) override
        {
            int ret;
            
            AuUInt32 done{}, read{};
            while (read != input)
            {
                AuUInt32 request = std::min(input, AuUInt32(ArraySize(din_)));
                if (this->_reader->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return std::make_pair(read, done);
                }
                read += 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 std::make_pair(read, 0);
                    }

                    auto have = ArraySize(dout_) - this->_ctx.avail_out;
                    done += have;

                    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 std::make_pair(read, done);
        }

    private:
        AuList<AuUInt8> _outbuffer;
        Aurora::IO::IStreamReader *_reader;
        bz_stream _ctx {};
        bool _init {};
        char dout_[4096];
        char din_[4096];

    };

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

        std::pair<AuUInt32, AuUInt32> Ingest(AuUInt32 input) override
        {
            AuUInt32 done {}, read {}, lastFrameSize {};
            std::shared_ptr<char> inFrame;
            std::shared_ptr<char> outFrames[2];
            bool outFrame {};

            while (read != input)
            {
                AuUInt16 frameSize;
                AuUInt32 request = sizeof(frameSize);

                if (this->_reader->Read(&frameSize, request) != IO::EStreamError::eErrorNone)
                {
                    return std::make_pair(read, done);
                }
                read += request;
                
                if ((lastFrameSize < frameSize) || (!outFrames[outFrame]))
                {
                    inFrame = std::shared_ptr<char>(new char[frameSize], std::default_delete<char[]>());
                    outFrames[outFrame] = std::shared_ptr<char>(new char[frameSize], std::default_delete<char[]>());
                    lastFrameSize = frameSize;
                }

                request = frameSize;
                if (this->_reader->Read(inFrame.get(), request) != IO::EStreamError::eErrorNone)
                {
                    return std::make_pair(read, done);
                }
                read += request;
                
                auto outPtr = outFrames[outFrame].get();

                auto bytes = LZ4_decompress_safe_continue(_lz4Stream, inFrame.get(), outPtr, frameSize, frameSize);
                if (bytes <= 0)
                {
                    return std::make_pair(read, 0);
                }

                done += bytes;

                this->_outbuffer.insert(this->_outbuffer.end(), 
                                        reinterpret_cast<const AuUInt8 *>(outPtr),
                                        reinterpret_cast<const AuUInt8 *>(outPtr) + bytes);
                
                outFrames[!outFrame] = std::move(outFrames[outFrame]);
                outFrame = !outFrame;
            }
            
            return std::make_pair(read, done);
        }

    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))
            {
                delete ret;
                ret = nullptr;
            }
        }

        return ret;
    }

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