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

    File: LZMACompressor.hpp
    Date: 2023-7-2
    Author: Reece
***/
#pragma once

#include <lzma.h>

namespace Aurora::Compression
{
    struct LZMADeflate : BaseStream
    {
        CompressInfo meta;
        lzma_stream stream;

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

        ~LZMADeflate()
        {
            if (this->init_)
            {
                lzma_end(&this->stream);
            }
        }

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

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

            lzma_ret ret {};

            if (meta.uCompressionLevel > 9 ||
                meta.uCompressionLevel < 0)
            {
                SysPushErrorArg();
                return false;
            }

            if (!meta.uCompressionLevel)
            {
                meta.uCompressionLevel = 6;
            }

            lzma_check check {};
            if (meta.bErrorCheck)
            {
                check = LZMA_CHECK_SHA256;
            }
            else
            {
                check = LZMA_CHECK_NONE;
            }

            if (meta.uThreads > 1)
            {
                lzma_mt options {};
                options.block_size = meta.uBlockSize;
                options.preset = meta.uCompressionLevel;
                options.check = check;
                options.threads = AuMin(AuHwInfo::GetCPUInfo().uThreads, meta.uThreads);
                ret = lzma_stream_encoder_mt(&this->stream, &options);

                if (ret != LZMA_OK)
                {
                    this->SetLastError(ret);
                    AuResetMember(this->pReader_);
                    return false;
                }
            }
            else
            {
                ret = lzma_easy_encoder(&this->stream, meta.uCompressionLevel, check);

                if (ret != LZMA_OK)
                {
                    this->SetLastError((int)ret);
                    SysPushErrorGeneric("LZMA Coder failure: {}", (int)ret);
                    AuResetMember(this->pReader_);
                    return false;
                }
            }

            this->init_ = true;

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

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

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

            while (read < input)
            {
                z_stream ctx_ {};

                read += IngestForInPointer<const uint8_t, size_t>(this->pReader_, this->stream.next_in, this->stream.avail_in, input - read, this);

                if (!this->stream.avail_in)
                {
                    return { read, done };
                }

                do
                {
                    auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
                    this->stream.avail_out = uMainDOutLength;
                    this->stream.next_out = (Bytef *)pMainDOut;

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

                    auto ret = lzma_code(&this->stream, LZMA_RUN);
                    if (ret != LZMA_OK &&
                        ret != LZMA_STREAM_END)
                    {
                        this->SetLastError(ret);
                        this->pReader_.reset();
                        return AuMakePair(read, 0);
                    }

                    auto have = uMainDOutLength - this->stream.avail_out;
                    done += have;

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

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

            return { read, done };
        }

        bool Flush() override
        {
            return RunFlush(LZMA_FULL_FLUSH);
        }

        bool Finish() override
        {
            return RunFlush(LZMA_FINISH);
        }

        bool RunFlush(lzma_action action)
        {
            if (!this->pReader_)
            {
                return false;
            }

            while (this->stream.avail_in)
            {
                do
                {
                    auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
                    this->stream.avail_out = uMainDOutLength;
                    this->stream.next_out = (Bytef *)pMainDOut;

                    auto ret = lzma_code(&this->stream, LZMA_RUN);
                    if (ret != LZMA_OK &&
                        ret != LZMA_STREAM_END)
                    {
                        this->SetLastError(int(ret));
                        this->pReader_.reset();
                        return false;
                    }

                    auto have = uMainDOutLength - this->stream.avail_out;

                    if (!Write2(reinterpret_cast<const AuUInt8 *>(this->dout_), have))
                    {
                        this->pReader_.reset();
                        this->SetLastError(0x69, "OOM");
                        return false;
                    }
                }
                while (this->stream.avail_out == 0);
            }

            if (!this->stream.avail_in)
            {
                do
                {
                    auto [pMainDOut, uMainDOutLength] = this->GetDOutPair();
                    this->stream.avail_out = uMainDOutLength;
                    this->stream.next_out = (Bytef *)pMainDOut;

                    auto ret = lzma_code(&this->stream, action);
                    if (ret != LZMA_OK &&
                        ret != LZMA_STREAM_END)
                    {
                        this->SetLastError(int(ret));
                        this->pReader_.reset();
                        return false;
                    }

                    auto have = uMainDOutLength - this->stream.avail_out;

                    if (!Write2(pMainDOut, have))
                    {
                        this->pReader_.reset();
                        this->SetLastError(0x69, "OOM");
                        return false;
                    }
                }
                while (this->stream.avail_out == 0);

                return true;
            }

            return true;
        }

    private:
        AuSPtr<IO::IStreamReader> pReader_;
        bool init_ {};
        unsigned char din_[kChunkSize];
        unsigned char dout_[kChunkSize];
    };
}