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

    File: LZ4Compressor.hpp
    Date: 2022-2-15
    Author: Reece
***/
#pragma once

#include "lz4.h"
#include "lz4frame.h"

namespace Aurora::Compression
{
    struct LZ4Deflate : public BaseStream
    {
        CompressionInfo meta;
        AuSPtr<char> bufferIn_;
        AuSPtr<char> bufferOut_;
        char* readPtr_;
        AuUInt32 bufferInAvail{};

        LZ4Deflate(const CompressionInfo &meta) : meta(meta), BaseStream(meta.internalStreamSize)
        {
            if (meta.bLZ4AutoFlush)
            {
                pref.autoFlush = true;
            }
        }

        ~LZ4Deflate()
        {
           
        }

        bool Init(const AuSPtr<IO::IStreamReader> &reader) override
        {
            this->reader_ = reader;

            auto err = LZ4F_createCompressionContext(&this->cctxPtr, LZ4F_getVersion());
            if (LZ4F_isError(err))
            {
                return {};
            }

            auto bufferSize = meta.internalStreamSize;

            this->bufferIn_ = AuSPtr<char>(new char[bufferSize], AuDefaultDeleter<char[]>());

            if (!this->bufferIn_)
            {
                return {};
            }

            this->bufferOut_ = AuSPtr<char>(new char[bufferSize], AuDefaultDeleter<char[]>());
            if (!this->bufferOut_)
            {
                return {};
            }

            this->readPtr_ = this->bufferIn_.get();
            SetPointer(this->readPtr_, bufferSize);
            return true;
        }

        bool AU_NOINLINE  BeginLZ4FrameIfNeeded()
        {
            if (bDead)
            {
                return false;
            }

            if (AuExchange(this->bHasEncodedFrame, true))
            {
                return true;
            }

            return DoLZ4Start();
        }

        bool AU_NOINLINE EndLZ4FrameIfNeeded()
        {
            if (!this->bHasEncodedFrame)
            {
                return true;
            }

            if (AuExchange(this->bHasFinished, true))
            {
                return true;
            }

            return DoLZ4End();
        }

        bool AU_NOINLINE  DoLZ4Start()
        {
            this->bytesRemInFrame = 0;
            this->bHasFinished = false;

            auto written = LZ4F_compressBegin(this->cctxPtr, this->bufferOut_.get(), meta.internalStreamSize, &this->pref);
            if (LZ4F_isError(written))
            {
                return false;
            }
            bytesRemInFrame = written;

            if (written)
            {
                if (!Write(this->bufferOut_.get(), written))
                {
                    return {};
                }
            }

            return true;
        }

        bool AU_NOINLINE DoLZ4End()
        {
            if (!this->bufferOut_)
            {
                return false;
            }

            auto startPtr = this->bufferOut_.get() + this->bytesRemInFrame;
            AuUInt32 bufferedBytes = LZ4F_compressEnd(cctxPtr, startPtr, meta.internalStreamSize - bytesRemInFrame, &options);
            if (LZ4F_isError(bufferedBytes))
            {
                this->bDead = true;
                return false;
            }

            this->bytesRemInFrame += bufferedBytes;
            this->bHasEncodedFrame = false;

            if (bufferedBytes)
            {
                if (!Write(startPtr, bufferedBytes))
                {
                    return {};
                }
            }

            return true;
        }

        AuStreamReadWrittenPair_t Ingest_s(AuUInt32 input) override
        {
            bool ret = true;
            AuUInt32 inputStat = 0, outputStat = 0;

            auto startLen = bytesRemInFrame;

            if (!BeginLZ4FrameIfNeeded())
            {
                return {};
            }

            auto bufferSize = meta.internalStreamSize;

            while (inputStat < input)
            {
                inputStat += IngestForInPointer<char, uInt>(this->reader_, this->readPtr_, this->bufferInAvail, input - inputStat);

                if (!this->bufferInAvail)
                {
                    return { inputStat, (startLen - bytesRemInFrame) };
                }

                size_t frameSPtr = this->bufferInAvail;
                size_t frameS2Ptr = bufferSize - bytesRemInFrame;
                
                auto startPtr = this->bufferOut_.get() + bytesRemInFrame;
                auto temp  = LZ4F_compressUpdate(this->cctxPtr, startPtr, frameS2Ptr, this->readPtr_, frameSPtr, &options);
                if (LZ4F_isError(temp))
                {
                    SysPushErrorGeneric("LZ4 internal stream size was too small. ingested too much data. must reset stream now");
                    bDead = true;
                    return {};
                }

                bytesRemInFrame += temp;

                this->readPtr_ += this->bufferInAvail;
                this->bufferInAvail = 0;

                if (temp)
                {
                    if (!Write(startPtr, temp))
                    {
                        return {};
                    }
                }

                outputStat += temp;
            }

            return AuMakePair(inputStat, (startLen - bytesRemInFrame));
        }

        bool Flush() override
        {
            return EndLZ4FrameIfNeeded();
        }

        bool Finish() override
        {
            return Flush();
        }

    private:

        bool bHasEncodedFrame{};
        bool bHasFinished{};
        bool bDead{};
        size_t bytesRemInFrame{};
        AuSPtr<IO::IStreamReader> reader_;
        LZ4F_cctx* cctxPtr;
        LZ4F_preferences_t pref{};
        LZ4F_compressOptions_t options{};
    };
}