/***
    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"
#include "lz4frame.h"

namespace Aurora::Compression
{
    bool BaseStream::ReadByProcessedN(void *buffer, AuUInt32 minimumInflated, AuStreamReadWrittenPair_t &pair, bool ingestUntilEOS)
    {
        AuUInt32 read {}, len {};

        if (ingestUntilEOS)
        {
            while (this->_outbuffer.RemainingBytes() < minimumInflated)
            {
                auto toRead = minimumInflated ? std::min(AuUInt32(4096), AuUInt32(minimumInflated - this->_outbuffer.RemainingBytes())) : 4096;
                if (Ingest(toRead).second == 0)
                {
                    if (!this->_outbuffer.RemainingBytes())
                    {
                        return false;
                    }

                    break;
                }

                read += toRead;
            }
        }

        len  = this->_outbuffer.Read(buffer, minimumInflated, buffer == nullptr);
        pair = {read, len};
        return len != 0;
    }

    bool BaseStream::ReadByProcessedN(void *buffer, AuUInt32 minimumInflated)
    {
        AuUInt32 read {}, len {};
        len = this->_outbuffer.Read(buffer, minimumInflated, buffer == nullptr);
        return len != 0;
    }

    bool BaseStream::GoBackByProcessedN(AuUInt32 offset)
    {
        return this->_outbuffer.ReaderTryGoBack(offset);
    }

    bool BaseStream::GoForwardByProcessedN(AuUInt32 offset)
    {
        return this->_outbuffer.ReaderTryGoForward(offset);
    }

    bool BaseStream::Write(const void *a, AuUInt32 length)
    {
        auto written = this->_outbuffer.Write(reinterpret_cast<const AuUInt8 *>(a),
                                              length);

        if (written != length)
        {
            auto increase = std::max(0, (int)length - (int)this->_outbuffer.RemainingWrite());
            increase += this->_outbuffer.length;

            if (increase > 64 * 1024 * 1024)
            {
                return false;
            }

            if (!this->_outbuffer.Resize(increase))
            {
                return false;
            }


            auto remaining = length - written;
            written = this->_outbuffer.Write(reinterpret_cast<const AuUInt8 *>(a) + written,
                                                remaining);
            if (written != remaining)
            {
                return false;
            }
        }

        return true;
    }
    
    AuUInt32 BaseStream::GetInternalBufferSize()
    {
        return this->_outbuffer.allocSize;
    }

    class ZSTDInflate : public BaseStream
    {
    public:
        DecompressInfo meta;
        ZSTDInflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
        {}

        ~ZSTDInflate()
        {
            if (auto dctx = std::exchange(dctx_, {}))
            {
                ZSTD_freeDCtx(dctx);
            }
        }

        bool Init(const AuSPtr<Aurora::IO::IStreamReader> &reader)
        {
            this->reader_ = reader;
            this->dctx_ = ZSTD_createDCtx();
            
            if (!this->dctx_)
            {
                SysPushErrorGen("Couldn't create decompressor");
                return false;
            }

            return true;
        }
        
        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            AuUInt32 length         = ZSTD_DStreamInSize();
            AuUInt32 outFrameLength = ZSTD_DStreamOutSize();
            AuUInt32 done{}, read{};

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

                    done += output.pos;
             
                    if (!Write(reinterpret_cast<const AuUInt8 *>(output.dst),
                        output.pos))
                    {
                        return {};
                    }
                }
            }
            
            return AuMakePair(read, done);
        }
  
    private:
      
        AuSPtr<Aurora::IO::IStreamReader> reader_;
        ZSTD_DCtx *dctx_;
        char din_[ZSTD_BLOCKSIZE_MAX + 3 /*ZSTD_BLOCKHEADERSIZE*/];
        char dout_[ZSTD_BLOCKSIZE_MAX];
        ZSTD_inBuffer input_;
    };

    class ZIPInflate : public BaseStream
    {
    public:
        DecompressInfo meta;

        ZIPInflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
        {}

        ~ZIPInflate()
        {
            if (this->init_)
            {
                inflateEnd(&this->ctx_);
            }
        }

        bool Init(const AuSPtr<Aurora::IO::IStreamReader> &reader)
        {
            this->reader_ = reader;
            
            auto ret = inflateInit(&this->ctx_);
            if (ret < Z_OK)
            {
                SysPushErrorMem("Error: {}", ret);
                return false;
            }

            this->init_ = true;
            return true;
        }

        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            int ret;

            AuUInt32 done{}, read{};
            
            while (read < input)
            {
                AuUInt32 request = std::min(input, AuUInt32(AuArraySize(din_)));
                if (this->reader_->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return AuMakePair(read, done);
                }

                read += request;

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

                do
                {
                    this->ctx_.avail_out = AuArraySize(dout_); // std::min(AuUInt32(AuArraySize(dout_)), AuUInt32(this->_outbuffer.RemainingWrite()));
                    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 AuMakePair(read, 0);
                    }

                    auto have = AuArraySize(dout_) - this->ctx_.avail_out;
                    done += have;

                    if (!Write(reinterpret_cast<const AuUInt8 *>(dout_),
                        have))
                    {
                        return {};
                    }

                } while (this->ctx_.avail_out == 0);
            }
            
            return AuMakePair(read, done);
        }

    private:

        AuSPtr<Aurora::IO::IStreamReader> reader_;
        z_stream ctx_ {};
        bool init_ {};
        unsigned char din_[4096];
        unsigned char dout_[4096];
    };

    class BZIPInflate : public BaseStream
    {
    public:
        DecompressInfo meta;

        BZIPInflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
        {}

        ~BZIPInflate()
        {
            if (this->init_)
            {
                BZ2_bzDecompressEnd(&this->ctx_);
            }
        }

        bool Init(const AuSPtr<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->init_ = true;
            return true;
        }
        
        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            int ret;
            
            AuUInt32 done{}, read{};
            while (read < input)
            {
                AuUInt32 request = std::min(input, AuUInt32(AuArraySize(din_)));
                if (this->reader_->Read(din_, request) != IO::EStreamError::eErrorNone)
                {
                    return AuMakePair(read, done);
                }
                read += request;

                this->ctx_.avail_in = request;
                this->ctx_.next_in = reinterpret_cast<char *>(din_);

                do
                {
                    this->ctx_.avail_out = AuArraySize(dout_);
                    this->ctx_.next_out  = dout_;

                    ret = BZ2_bzDecompress(&this->ctx_);
                    if (ret < Z_OK)
                    {
                        SysPushErrorIO("Error: {}", ret);
                        return AuMakePair(read, 0);
                    }

                    auto have = AuArraySize(dout_) - this->ctx_.avail_out;
                    done += have;

                    if (!Write(reinterpret_cast<const AuUInt8 *>(dout_),
                        have))
                    {
                        return {};
                    }

                } while (this->ctx_.avail_out == 0);
            }

            return AuMakePair(read, done);
        }

    private:

        AuSPtr<Aurora::IO::IStreamReader> reader_;
        bz_stream ctx_ {};
        bool init_ {};
        char dout_[4096];
        bool userBound_ {};
        char din_[4096];
    };

    class LZ4Inflate : public BaseStream
    {
    public:

        DecompressInfo meta;

        LZ4Inflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
        {}

        ~LZ4Inflate()
        {
            if (lz4Stream_)
            {
                LZ4F_freeDecompressionContext(lz4Stream_);
            }
        }

        bool Init(const AuSPtr<Aurora::IO::IStreamReader> &reader)
        {
            this->reader_ = reader;
            
            auto err = LZ4F_createDecompressionContext(&lz4Stream_, LZ4F_getVersion());
            if (LZ4F_isError(err))
            {
                return {};
            }

            return true;
        }


        AuStreamReadWrittenPair_t Ingest(AuUInt32 input) override
        {
            bool ret = true;
            LZ4F_dctx *dctxPtr;
            AuUInt32 inputStat = 0, outputStat = 0;
            size_t bytesRemInFrame {};
            LZ4F_decompressOptions_t opts {};

            auto bufferSize = meta.internalStreamSize / 2;

            auto bufferIn  = AuSPtr<char>(new char[bufferSize], std::default_delete<char[]>());
            auto bufferOut = AuSPtr<char>(new char[bufferSize], std::default_delete<char[]>());

            while (inputStat < input)
            {
                auto frameSize = bytesRemInFrame ? bytesRemInFrame : LZ4F_MIN_SIZE_TO_KNOW_HEADER_LENGTH;
                auto min = frameSize;

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

                if (frameSize)
                {
                    AuUInt32 request = frameSize;
                    if ((input < (inputStat + request)) ||
                        (this->reader_->Read(bufferIn.get(), request) != IO::EStreamError::eErrorNone) ||
                        (request != frameSize))
                    {
                        ret = request == 0 && inputStat;
                        break;
                    }
                    inputStat += frameSize;
                }

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

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

                    if (frameS2Ptr)
                    {
                        if (!Write(bufferOut.get(), frameS2Ptr))
                        {
                            ret = false;
                            break;
                        }
                    }

                    outputStat += frameS2Ptr;
                }
            }


            if (!ret)
            {
                return {};
            }

            return  AuMakePair(inputStat, outputStat);
        }

    private:

        AuSPtr<Aurora::IO::IStreamReader> reader_;
        LZ4F_dctx* lz4Stream_ {};
    };

    AUKN_SYM ICompressionStream *DecompressorNew(const AuSPtr<Aurora::IO::IStreamReader> &reader, const DecompressInfo &ref)
    {
        DecompressInfo info = ref;
        BaseStream * ret{};

        if (!info.internalStreamSize)
        {
            info.internalStreamSize = 1024 * 64 * 2;
        }

        switch (info.alg)
        {
        case ECompresionType::eZSTD:
            ret = new ZSTDInflate(info);
            break;
        case ECompresionType::eBZIP2:
            ret = new BZIPInflate(info);
            break;
        case ECompresionType::eLZ4:
            ret = new LZ4Inflate(info);
            break;
        case ECompresionType::eDeflate:
            ret = new ZIPInflate(info);
            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);
    }
}