AuroraRuntime/Source/IO/Protocol/ProtocolStack.cpp

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

    File: ProtocolStack.cpp
    Date: 2022-8-24
    Author: Reece
***/
#include "Protocol.hpp"
#include "ProtocolStack.hpp"
#include "ProtocolPiece.hpp"
#include "IProtocolNext.hpp"
#include <Source/IO/IOPipeProcessor.hpp>

namespace Aurora::IO::Protocol
{
    ProtocolStack::~ProtocolStack()
    {
    }

    bool ProtocolStack::AppendInterceptor(const AuSPtr<IProtocolInterceptor> &pInterceptor,
                                          AuUInt uOutputBufferSize)
    {
        return this->AddInterceptorWhere(false, pInterceptor, uOutputBufferSize);
    }

    bool ProtocolStack::AppendInterceptorEx(const AuSPtr<IProtocolInterceptorEx> &pInterceptorEx,
                                            AuUInt uOutputBufferSize)
    {
        return this->AddInterceptorWhereEx(false, pInterceptorEx, uOutputBufferSize);
    }

    bool ProtocolStack::PrependInterceptor(const AuSPtr<IProtocolInterceptor> &pInterceptor,
                                           AuUInt uOutputBufferSize)
    {
        return this->AddInterceptorWhere(true, pInterceptor, uOutputBufferSize);
    }

    bool ProtocolStack::PrependInterceptorEx(const AuSPtr<IProtocolInterceptorEx> &pInterceptorEx,
                                             AuUInt uOutputBufferSize)
    {
        return this->AddInterceptorWhereEx(true, pInterceptorEx, uOutputBufferSize);
    }

    bool ProtocolStack::AddInterceptorWhere(bool prepend,
                                            const AuSPtr<IProtocolInterceptor> &pInterceptor,
                                            AuUInt uOutputBufferSize)
    {
        if (this->bWrittenEnd)
        {
            return false;
        }

        if (!uOutputBufferSize)
        {
            uOutputBufferSize = 64 * 1024;
        }

        auto pNew = AuMakeShared<ProtocolPiece>();
        if (!pNew)
        {
            SysPushErrorNet("Out of memory");
            return false;
        }

        pNew->outputBuffer = AuByteBuffer(uOutputBufferSize, true, false);
        if (!pNew->outputBuffer.IsValid())
        {
            SysPushErrorNet("Out of memory");
            return false;
        }

        // Circular ref
        pNew->pOuputWriter = AuMakeShared<AuIO::Buffered::BlobWriter>(AuSPtr<Memory::ByteBuffer>(pNew, &pNew->outputBuffer));

        struct StreamWrapper : IStreamWriter, IProtocolNext, AuEnableSharedFromThis<StreamWrapper>
        {
            ProtocolStack *pStack;
            AuWPtr<IProtocolInterceptor> pInterceptor;
            AuWPtr<ProtocolPiece> pParent;

            EStreamError IsOpen() override
            {
                return EStreamError::eErrorNone;
            }

            EStreamError Write(const Memory::MemoryViewStreamRead &parameters) override
            {
                return pStack->DoTick(AuMakeShared<AuByteBuffer>(parameters), pParent.lock()) ?
                    EStreamError::eErrorNone :
                    EStreamError::eErrorStreamInterrupted;
            }

            void Close() override
            {

            };

            void Flush() override
            {

            };

            virtual AuSPtr<Memory::ByteBuffer> GetOutputBuffer() override
            {
                return {};
            }

            virtual AuSPtr<IStreamWriter> GetStreamWriter() override
            {
                return AuSharedFromThis();
            }
        };

        auto pWrapper = AuMakeShared<StreamWrapper>();
        if (!pWrapper)
        {
            return false;
        }

        pWrapper->pInterceptor = pInterceptor;
        pWrapper->pStack = this;
        pWrapper->pParent = pNew;

        pNew->pWriteInteface = pWrapper;
        pNew->pParent = this;
        pNew->pInterceptor = pInterceptor;

        if (prepend)
        {
            if (this->pBottomPiece)
            {
                pNew->pNext = this->pBottomPiece;
            }

            this->pBottomPiece = pNew;
        }
        else
        {
            this->pTopPiece = pNew;

            if (!this->pBottomPiece)
            {
                this->pBottomPiece = pNew;
            }
        }

        return true;
    }
    
    bool ProtocolStack::AddInterceptorWhereEx(bool prepend,
                                              const AuSPtr<IProtocolInterceptorEx> &pInterceptor,
                                              AuUInt uOutputBufferSize)
    {
        if (this->bWrittenEnd)
        {
            return false;
        }

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

        if (!uOutputBufferSize)
        {
            uOutputBufferSize = 64 * 1024;
        }

        auto pNew = AuMakeShared<ProtocolPiece>();
        if (!pNew)
        {
            SysPushErrorNet("Out of memory");
            return false;
        }

        pNew->outputBuffer = AuByteBuffer(uOutputBufferSize, true, false);
        if (!pNew->outputBuffer.IsValid())
        {
            SysPushErrorNet("Out of memory");
            return false;
        }

        // Circular ref
        pNew->pOuputWriter = AuMakeShared<AuIO::Buffered::BlobWriter>(AuSPtr<Memory::ByteBuffer>(pNew, &pNew->outputBuffer));

        struct StreamWrapper : IStreamWriter, IProtocolNext, AuEnableSharedFromThis<StreamWrapper>
        {
            ProtocolStack *pStack;
            AuWPtr<IProtocolInterceptorEx> pInterceptor;
            AuWPtr<ProtocolPiece> pParent;

            EStreamError IsOpen() override
            {
                return EStreamError::eErrorNone;
            }

            EStreamError Write(const Memory::MemoryViewStreamRead &parameters) override
            {
                return pStack->DoTick(AuMakeShared<AuByteBuffer>(parameters), pParent.lock()) ?
                        EStreamError::eErrorNone :
                        EStreamError::eErrorStreamInterrupted;
            }

            void Close() override
            {

            };

            void Flush() override
            {

            };

            virtual AuSPtr<Memory::ByteBuffer> GetOutputBuffer() override
            {
                return {};
            }

            virtual AuSPtr<IStreamWriter> GetStreamWriter() override
            {
                return AuSharedFromThis();
            }
        };

        auto pWrapper = AuMakeShared<StreamWrapper>();
        if (!pWrapper)
        {
            return false;
        }

        pWrapper->pInterceptor = pInterceptor;
        pWrapper->pStack = this;
        pWrapper->pParent = pNew;

        pNew->pWriteInteface = pWrapper;
        pNew->pInterceptorEx = pInterceptor;
        pNew->pParent = this;

        if (prepend)
        {
            if (this->pBottomPiece)
            {
                pNew->pNext = this->pBottomPiece;
            }

            this->pBottomPiece = pNew;
        }
        else
        {
            this->pTopPiece = pNew;

            if (!this->pBottomPiece)
            {
                this->pBottomPiece = pNew;
            }
        }

        return true;
    }

    bool ProtocolStack::AddEndInterceptor(const AuSPtr<IProtocolInterceptorEx> &pInterceptor)
    {
        if (this->bWrittenEnd)
        {
            return false;
        }

        if (!this->pSourceBufer)
        {
            return false;
        }
   
        auto pNew = AuMakeShared<ProtocolPiece>();
        if (!pNew)
        {
            SysPushErrorNet("Out of memory");
            return false;
        }

        struct StreamWrapper : IStreamWriter, IProtocolNext, AuEnableSharedFromThis<StreamWrapper>
        {
            ProtocolStack *pStack;
            AuWPtr<IProtocolInterceptorEx> pInterceptor;
            AuWPtr<ProtocolPiece> pParent;

            EStreamError IsOpen() override
            {
                return EStreamError::eErrorNone;
            }

            EStreamError Write(const Memory::MemoryViewStreamRead &parameters) override
            {
                return pStack->DoTick(AuMakeShared<AuByteBuffer>(parameters), pParent.lock()) ?
                    EStreamError::eErrorNone :
                    EStreamError::eErrorStreamInterrupted;
            }

            void Close() override
            {

            };

            void Flush() override
            {

            };

            virtual AuSPtr<Memory::ByteBuffer> GetOutputBuffer() override
            {
                return {};
            }

            virtual AuSPtr<IStreamWriter> GetStreamWriter() override
            {
                return AuSharedFromThis();
            }
        };

        auto pWrapper = AuMakeShared<StreamWrapper>();
        if (!pWrapper)
        {
            return false;
        }

        pWrapper->pInterceptor = pInterceptor;
        pWrapper->pStack = this;
        pWrapper->pParent = pNew;

        pNew->pWriteInteface = pWrapper;
        pNew->pInterceptorEx = pInterceptor;
        pNew->pParent = this;

        this->pTopPiece = pNew;

        if (!this->pBottomPiece)
        {
            this->pBottomPiece = pNew;
        }

        this->bWrittenEnd = true;

        return true;
    }

    void ProtocolStack::Destroy()
    {
        if (this->bOwnsSource)
        {
            if (this->pSourceBufer)
            {
                this->pSourceBufer->Reset();
                this->pSourceBufer.reset();
            }
        }

        auto pItr = this->pBottomPiece;
        while (pItr)
        {
            auto pCur = pItr;
            pItr = pCur->pNext;
            pCur->pNext.reset();
            pCur->pInterceptor.reset();
            pCur->pInterceptorEx.reset();
            pCur->pOuputWriter.reset();
            pCur->pWriteInteface.reset();
        }

        this->pBottomPiece.reset();
        this->pTopPiece.reset();
    }

    AuSPtr<IStreamWriter> ProtocolStack::AsStreamWriter()
    {
        if (!this->bOwnsSource)
        {
            return {};
        }

        return AuMakeShared<AuIO::Buffered::BlobWriter>(this->pSourceBufer);
    }

    AuSPtr<Memory::ByteBuffer> ProtocolStack::AsWritableByteBuffer()
    {
        if (!this->bOwnsSource)
        {
            return {};
        }

        return this->pSourceBufer;
    }

    AuSPtr<IStreamReader> ProtocolStack::AsStreamReader()
    {
        auto pBuffer = AsReadableByteBuffer();
        if (!pBuffer)
        {
            return {};
        }

        return AuMakeShared<AuIO::Buffered::BlobReader>(pBuffer);
    }

    AuSPtr<Memory::ByteBuffer> ProtocolStack::AsReadableByteBuffer()
    {
        if (!this->pTopPiece)
        {
            return this->pSourceBufer;
        }

        if (this->pTopPiece->outputBuffer.IsEmpty())
        {
            return {};
        }

        return AuSPtr<AuByteBuffer>(this->pTopPiece, &this->pTopPiece->outputBuffer);
    }

    void ProtocolStack::DoTick()
    {
        if (!this->pSourceBufer)
        {
            return;
        }

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

        this->DoTick(this->pSourceBufer, {});
    }

    bool ProtocolStack::DoTick(const AuSPtr<AuByteBuffer> &pRead, const AuSPtr<ProtocolPiece> &pPiece)
    {
        auto pCurrent = pPiece ? pPiece : this->pBottomPiece;

        if (pCurrent->pInterceptorEx)
        {
            auto pNextStream = ((pPiece == this->pTopPiece) && (this->pDrainBuffer)) ?
                (this->pDrainBuffer) :
                (!pCurrent->outputBuffer.IsEmpty() ? AuSPtr<AuByteBuffer>(pCurrent, &pCurrent->outputBuffer) : AuSPtr<AuByteBuffer> {});

            auto pOldHead = pNextStream ? pNextStream->readPtr : nullptr;

            if (!pCurrent->pInterceptorEx->OnDataAvailable(pRead, pNextStream))
            {
                if (pNextStream)
                {
                    pNextStream->readPtr = pOldHead;
                }
                return false;
            }

            if (!pNextStream)
            {
                return true;
            }

            if (auto pNext = pCurrent->pNext)
            {
                auto pOldHead = pNextStream->readPtr;

                bool bDontCareAboutSubTicks = this->DoTick(pNextStream, pNext);

                if (!bDontCareAboutSubTicks)
                {
                    pNextStream->readPtr = pOldHead;
                }
            }

            return true;
        }
        else
        {
            AuUInt8 *pBase {};
            AuUInt uCount {};

            if (pRead->writePtr >= pRead->readPtr)
            {
                uCount = pRead->writePtr - pRead->readPtr;
                pBase  = pRead->readPtr;
            }
            else if (pRead->flagCircular)
            {
                uCount = (pRead->base + pRead->length) - pRead->readPtr;
                pBase  = pRead->readPtr;
            }

            auto pNextStream = pCurrent->ToNextWriter();

            if (!pCurrent->pInterceptor->OnDataAvailable({ pBase, uCount }, pNextStream))
            {
                return false;
            }

            return true;
        }
    }

    AUKN_SYM AuSPtr<IProtocolStack> NewBufferedProtocolStack(AuUInt uLength)
    {
        if (!uLength)
        {
            SysPushErrorArg();
            return {};
        }

        auto pStack = AuMakeShared<ProtocolStack>();
        if (!pStack)
        {
            SysPushErrorMem();
            return {};
        }

        auto pBuffer = AuMakeShared<AuByteBuffer>(uLength, true);
        if (!pBuffer)
        {
            SysPushErrorMem();
            return {};
        }

        if (!pBuffer->IsValid())
        {
            return {};
        }

        pStack->bOwnsSource = true;
        pStack->pSourceBufer = pBuffer;

        return pStack;
    }

    AUKN_SYM AuSPtr<IProtocolStack> NewProtocolStackFromPipe(const AuSPtr<IIOPipeWork> &pWork)
    {
        if (!pWork)
        {
            SysPushErrorArg();
            return {};
        }

        auto pStack = AuMakeShared<ProtocolStack>();
        if (!pStack)
        {
            SysPushErrorMem();
            return {};
        }

        pStack->pSourceBufer = AuSPtr<AuByteBuffer>(pWork, AuStaticCast<IOPipeWork>(pWork)->GetBuffer());

        return pStack;
    }
}