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

    File: BrotliCompressor.hpp
    Date: 2023-3-12
    Author: Reece
***/
#pragma once

#include "brotli/encode.h"

namespace Aurora::Compression
{
    struct BrotliDeflate : BaseStream
    {
        CompressInfo meta;

        BrotliDeflate(const CompressInfo &meta) : meta(meta), BaseStream(meta.uInternalStreamSize)
        { }

        ~BrotliDeflate()
        {
            if (this->pState)
            {
                BrotliEncoderDestroyInstance(this->pState);
            }
        }

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

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

            if (meta.uCompressionLevel > BROTLI_MAX_QUALITY)
            {
                SysPushErrorArg("Invalid brotli compression level");
                return false;
            }

            this->pState = BrotliEncoderCreateInstance(nullptr, nullptr, nullptr);
            if (!this->pState)
            {
                SysPushErrorMem("No brotli encoder");
                return false;
            }

            if (meta.uBlockSize)
            {
                if (!BrotliEncoderSetParameter(this->pState, BrotliEncoderParameter::BROTLI_PARAM_LGBLOCK, meta.uBlockSize))
                {
                    SysPushErrorArg("Compressor couldn't set block size");
                    return false;
                }
            }

            if (meta.uOptQuality)
            {
                if (meta.uOptQuality.value() <= BROTLI_MAX_QUALITY &&
                    meta.uOptQuality.value() >= BROTLI_MIN_QUALITY)
                {
                    if (!BrotliEncoderSetParameter(this->pState, BrotliEncoderParameter::BROTLI_PARAM_QUALITY, meta.uOptQuality.value()))
                    {
                        SysPushErrorArg("Compressor couldn't set quality");
                        return false;
                    }
                }
                else
                {
                    SysPushErrorArg("Compressor couldn't set quality");
                    return false;
                }
            }
            else
            {
                if (!BrotliEncoderSetParameter(this->pState, BrotliEncoderParameter::BROTLI_PARAM_QUALITY, meta.uCompressionLevel))
                {
                    SysPushErrorArg("Compressor couldn't set quality from compression level");
                    return false;
                }
            }

            if (meta.uOptWindowBits)
            {
                if (meta.uOptWindowBits.value() <= BROTLI_MAX_INPUT_BLOCK_BITS &&
                    meta.uOptWindowBits.value() >= BROTLI_MIN_INPUT_BLOCK_BITS)
                {
                    if (!BrotliEncoderSetParameter(this->pState, BrotliEncoderParameter::BROTLI_PARAM_LGWIN, meta.uOptWindowBits.value()))
                    {
                        SysPushErrorArg("Compressor couldn't set window bits");
                        return false;
                    }
                }
                else
                {
                    SysPushErrorArg("Compressor couldn't set window bits");
                    return false;
                }
            }

            for (const auto &[a, b] : meta.options)
            {
                if (!BrotliEncoderSetParameter(this->pState, (BrotliEncoderParameter)a, b))
                {
                    SysPushErrorArg("Compressor argument assignment {} = {}", a, b);
                    return false;
                }
            }

            this->InitByDesc(this->meta);
            this->SetArray(this->din_);
            this->SetOutArray(this->dout_);
            return true;
        }

        AuStreamReadWrittenPair_t Ingest_s(AuUInt32 input) override
        {
            AuUInt32 done {}, read {};

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

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

            while (read < input)
            {
                read += IngestForInPointer<uint8_t, size_t>(this->pReader_, this->pInBuffer, this->uAvailIn, input - read, this);

                if (!this->uAvailIn)
                {
                    return { read, done };
                }

                auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
                size_t outNext {};
                uint8_t *pOut {};
                do
                {
                    outNext = uMainDOutLength;
                    pOut = (uint8_t *)pMainDOut;
        
                    auto ret = BrotliEncoderCompressStream(this->pState, BrotliEncoderOperation::BROTLI_OPERATION_PROCESS, &this->uAvailIn, (const uint8_t **) & this->pInBuffer, &outNext, &pOut, nullptr);
                    if (ret == BROTLI_FALSE)
                    {
                        SysPushErrorIO("Brotli Error");
                        this->pReader_.reset();
                        return AuMakePair(read, 0);
                    }

                    auto have = uMainDOutLength - outNext;
                    done += have;

                    if (!Write2(reinterpret_cast<const AuUInt8 *>(pMainDOut), have))
                    {
                        this->pReader_.reset();
                        this->SetLastError(0x69, "OOM");
                        return AuMakePair(read, 0);
                    }
                }
                while (outNext == 0 || BrotliEncoderHasMoreOutput(this->pState));
            }

            return { read, done };
        }

        bool Flush() override
        {
            return RunFlush(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH);
        }

        bool Finish() override
        {
            return RunFlush(BrotliEncoderOperation::BROTLI_OPERATION_FINISH);
        }

        bool RunFlush(BrotliEncoderOperation type)
        {
            if (!this->pReader_)
            {
                return false;
            }
            if (!this->pState)
            {
                return false;
            }

            AuUInt read {};
            {
                size_t outNext {};
                uint8_t *pOut {};
                while (outNext == 0 || BrotliEncoderHasMoreOutput(this->pState))
                {
                    auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();

                    outNext = uMainDOutLength;
                    pOut = (uint8_t *)pMainDOut;

                    auto ret = BrotliEncoderCompressStream(this->pState, type, &this->uAvailIn, (const uint8_t **)&this->pInBuffer, &outNext, &pOut, nullptr);
                    if (ret == BROTLI_FALSE)
                    {
                        SysPushErrorIO("Brotli Error");
                        this->pReader_.reset();
                        return false;
                    }

                    auto have = uMainDOutLength - outNext;
                    if (!Write2(reinterpret_cast<const AuUInt8 *>(pMainDOut), have))
                    {
                        this->pReader_.reset();
                        this->SetLastError(0x69, "OOM");
                        return false;
                    }
                }
            }

            return true;
        }

    private:
        AuSPtr<IO::IStreamReader> pReader_;
        unsigned char din_[kChunkSize];
        unsigned char dout_[kChunkSize];
        uint8_t *pInBuffer;
        size_t uAvailIn {};
        BrotliEncoderState *pState {};
    };
}