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

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

#include <lzma.h>

namespace Aurora::Compression
{
    struct LZMAInflate : BaseStream
    {
        DecompressInfo meta;
        lzma_stream stream;

        LZMAInflate(const DecompressInfo &meta) : meta(meta), BaseStream(meta.uInternalStreamSize)
        {

        }

        ~LZMAInflate()
        {
            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 {};

            AuUInt32 uFlags {};
            if (meta.bErrorCheck)
            {
                uFlags = 0;
            }
            else
            {
                uFlags = LZMA_IGNORE_CHECK | LZMA_TELL_NO_CHECK;
            }

            if (meta.uThreads > 1)
            {
                lzma_mt options {};
                options.threads = AuMin(AuHwInfo::GetCPUInfo().uThreads, meta.uThreads);
                options.flags = uFlags;
                ret = lzma_stream_decoder_mt(&this->stream, &options);

                if (ret != LZMA_OK)
                {
                    this->SetLastError(ret);
                    AuResetMember(this->pReader_);
                    return false;
                }
            }
            else
            {
                ret = lzma_stream_decoder(&this->stream, UINT64_MAX, uFlags);

                if (ret != LZMA_OK)
                {
                    this->SetLastError(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 };
        }

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