/***
    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 : BaseStream
    {
        CompressInfo meta;

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

        ~LZ4Deflate()
        {
           
        }

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

            if (!this->IsValid())
            {
                SysPushErrorMem();
                return false;
            }

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

            auto bufferSize = meta.uInternalStreamSize;

            this->pBufferIn_ = AuMakeSharedArray<char>(bufferSize);
            if (!this->pBufferIn_)
            {
                return {};
            }

            this->pBufferOut_ = AuMakeSharedArray<char>(bufferSize);
            if (!this->pBufferOut_)
            {
                return {};
            }

            this->pReadPtr_ = this->pBufferIn_.get();
            SetPointer(this->pReadPtr_, bufferSize);
            return true;
        }

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

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

            return DoLZ4Start();
        }

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

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

            return DoLZ4End();
        }

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

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

            if (written)
            {
                if (!Write(this->pBufferOut_.get(), written))
                {
                    this->bDead = true;
                    return {};
                }
            }

            return true;
        }

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

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

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

            if (bufferedBytes)
            {
                if (!Write(startPtr, bufferedBytes))
                {
                    this->bDead = true;
                    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.uInternalStreamSize;

            while (inputStat < input)
            {
                inputStat += IngestForInPointer<char, uInt>(this->pReader_, this->pReadPtr_, this->uBufferInAvail_, input - inputStat);

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

                size_t frameSPtr = this->uBufferInAvail_;
                size_t frameS2Ptr = bufferSize - bytesRemInFrame;
                
                auto startPtr = this->pBufferOut_.get() + bytesRemInFrame;
                auto temp  = LZ4F_compressUpdate(this->cctxPtr, startPtr, frameS2Ptr, this->pReadPtr_, 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->pReadPtr_ += this->uBufferInAvail_;
                this->uBufferInAvail_ = 0;

                if (temp)
                {
                    if (!Write(startPtr, temp))
                    {
                        this->bDead = true;
                        return {};
                    }
                }

                outputStat += temp;
            }

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

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

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

    private:

        AuSPtr<char> pBufferIn_;
        AuSPtr<char> pBufferOut_;
        char *pReadPtr_;
        AuUInt32 uBufferInAvail_ {};
        bool bHasEncodedFrame{};
        bool bHasFinished{};
        bool bDead{};
        size_t bytesRemInFrame{};
        AuSPtr<IO::IStreamReader> pReader_;
        LZ4F_cctx* cctxPtr;
        LZ4F_preferences_t pref{};
        LZ4F_compressOptions_t options{};
    };
}