#include "impl.h"

#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/thread/threadedqueue.h>

using namespace NCache;
using namespace NSrvKernel;

THolder<IIoInput> NCache::Load(TReqCacher* cacher, IMeta* meta, TThreadedQueue* queue, TInstant deadline, const TString& id) {
    class TLoad
        : public TThreadedQueue::ICallback
        , public TNonCopyable
        , public THolder<IInputStream>
    {
    public:
        TLoad(TReqCacher* cacher, IMeta* meta, const TString& id) noexcept
            : Cacher_(cacher)
            , Meta_(meta)
            , Id_(id)
        {}

        const std::exception_ptr& Exception() const noexcept {
            return Exception_;
        }

    private:
        void DoRun() noexcept override try {
            if (Awaited()) {
                Cacher_->Load(Id_).Swap(*this);

                if (this->Get()) {
                    Meta_->Load(this->Get());
                }
            }
        } catch (...) {
            this->Destroy();
            Exception_ = std::current_exception();
        }

    private:
        TReqCacher* const Cacher_{ nullptr };
        IMeta* const Meta_{ nullptr };
        const TString Id_;
        std::exception_ptr Exception_;
    };

    class TInput: public IIoInput {
    private:
        class TRecver
            : public TThreadedQueue::ICallback
            , public TChunkList
        {
        private:
            enum {
                PoolSize = 10
            };

        public:
            TRecver(THolder<IInputStream> stream) noexcept
                : Stream_(std::move(stream))
            {}

            void Prepare() noexcept {
                for (size_t size = Chunks_.ChunksCount(); size < PoolSize; ++size) {
                    Chunks_.Push(NewChunkReserve());
                }
            }

            THolder<IInputStream> ReleaseStream() noexcept {
                return std::move(Stream_);
            }

            const std::exception_ptr& Exception() const noexcept {
                return Exception_;
            }

        private:
            void DoRun() noexcept override try {
                while (!Chunks_.Empty() && Awaited()) {
                    auto chunk = Chunks_.PopFront();

                    const size_t readed = Stream_->Read(chunk->Data(), chunk->Length());

                    if (!readed) {
                        Chunks_.PushFront(std::move(chunk));
                        break;
                    }

                    chunk->Shrink(readed);
                    PushBack(std::move(chunk));
                }
            } catch (...) {
                Exception_ = std::current_exception();
            }

        private:
            THolder<IInputStream> Stream_;
            TChunkList Chunks_;
            std::exception_ptr Exception_;
        };

    public:
        TInput(TThreadedQueue* queue, THolder<IInputStream> stream) noexcept
            : Queue_(queue)
            , Recver_(MakeHolder<TRecver>(std::move(stream)))
        {}

    private:
        TRecver* Recver() const noexcept {
            return static_cast<TRecver*>(Recver_.Get());
        }

    private:
        TError DoRecv(TChunkList& lst, TInstant) noexcept override {
            if (Y_UNLIKELY(!Recver_)) {
                return {};
            }

            Recver()->Swap(lst);
            Recver()->Prepare();
            Recver_ = THolder<TThreadedQueue::ICallback>(Queue_->Run(Recver_.Release(), TInstant::Max()).Release());

            if (Y_UNLIKELY(!Recver_)) {
                return {};
            }

            try {
                try {
                    if (Recver()->Exception()) {
                        std::rethrow_exception(Recver()->Exception());
                    }

                    Recver()->Swap(lst);
                    return {};
                } catch (...) {
                    Recver_.Destroy();
                    throw;
                }
            } Y_TRY_STORE(yexception);
        }

    private:
        TThreadedQueue* const Queue_ = nullptr;

        THolder<TThreadedQueue::ICallback> Recver_;
    };

    const THolder<TThreadedQueue::ICallback> callback(
        queue->Run<TThreadedQueue::ICallback>(new TLoad(cacher, meta, id), deadline).Release());

    if (!callback) {
        return nullptr;
    }

    TLoad* const load = static_cast<TLoad*>(callback.Get());

    if (load->Exception()) {
        std::rethrow_exception(load->Exception());
    }

    if (!*load) {
        return nullptr;
    }

    return MakeHolder<TInput>(queue, std::move(*load));
}

THolder<IIoOutput> NCache::Store(TReqCacher* cacher, IMeta* meta, TThreadedQueue* queue, TInstant deadline, const TString& id) {
    class TStore
        : public TThreadedQueue::ICallback
        , public TNonCopyable
        , public THolder<IOutputStream>
    {
    public:
        TStore(TReqCacher* cacher, IMeta* meta, const TString& id) noexcept
            : Cacher_(cacher)
            , Meta_(meta)
            , Id_(id)
        {}

        const std::exception_ptr& Exception() const noexcept {
            return Exception_;
        }

    private:
        void DoRun() noexcept override try {
            if (Awaited()) {
                Cacher_->Store(Id_).Swap(*this);

                if (this->Get()) {
                    Meta_->Store(this->Get());
                }
            }
        } catch (...) {
            this->Destroy();
            Exception_ = std::current_exception();
        }

    private:
        TReqCacher* const Cacher_{ nullptr };
        IMeta* const Meta_{ nullptr };
        const TString Id_;
        std::exception_ptr Exception_;
    };

    class TOutput: public IIoOutput {
    private:
        class TSender
            : public TThreadedQueue::ICallback
            , public TChunkList
        {
        public:
            TSender(THolder<IOutputStream> stream) noexcept
                : Stream_(std::move(stream))
            {}

            void Cleanup() noexcept {
                Chunks_.Clear();
            }

            THolder<IOutputStream> ReleaseStream() noexcept {
                return std::move(Stream_);
            }

            const std::exception_ptr& Exception() const noexcept {
                return Exception_;
            }

        private:
            void DoRun() noexcept override try {
                while (!Empty() && Awaited()) {
                    auto chunk = PopFront();

                    Stream_->Write(chunk->Data(), chunk->Length());

                    Chunks_.PushBack(std::move(chunk));
                }
            } catch (...) {
                Exception_ = std::current_exception();
            }

        private:
            THolder<IOutputStream> Stream_;
            TChunkList Chunks_;
            std::exception_ptr Exception_;
        };

    private:
        class TCommiter
            : public TThreadedQueue::ICallback
            , public THolder<IOutputStream>
        {
        public:
            TCommiter(TReqCacher* cacher, THolder<IOutputStream> stream) noexcept
                : THolder<IOutputStream>(std::move(stream))
                , Cacher_(cacher)
            {}

            const std::exception_ptr& Exception() const noexcept {
                return Exception_;
            }

        private:
            void DoRun() noexcept override try {
                Cacher_->Commit(std::move(*this));
            } catch (...) {
                Exception_ = std::current_exception();
            }

        private:
            TReqCacher* const Cacher_{ nullptr };
            std::exception_ptr Exception_;
        };

        class TDiscarder
            : public TThreadedQueue::ICallback
            , public THolder<IOutputStream>
        {
        public:
            TDiscarder(TReqCacher* cacher, THolder<IOutputStream> stream) noexcept
                : THolder<IOutputStream>(std::move(stream))
                , Cacher_(cacher)
            {}

            const std::exception_ptr& Exception() const noexcept {
                return Exception_;
            }

        private:
            void DoRun() noexcept override try {
                Cacher_->Discard(std::move(*this));
            } catch (...) {
                Exception_ = std::current_exception();
            }

        private:
            TReqCacher* const Cacher_{ nullptr };
            std::exception_ptr Exception_;
        };

    public:
        TOutput(TReqCacher* cacher, TThreadedQueue* queue, THolder<IOutputStream> stream)
            : Queue_(queue)
            , Sender_(MakeHolder<TSender>(std::move(stream)))
            , Commiter_(MakeHolder<TCommiter>(cacher, nullptr))
            // , Discarder_(new TDiscarder(cacher, 0))
        {}

    private:
        TSender* Sender() const noexcept {
            return static_cast<TSender*>(Sender_.Get());
        }

        TCommiter* Commiter() const noexcept {
            return static_cast<TCommiter*>(Commiter_.Get());
        }

        TDiscarder* Discarder() const noexcept {
            return static_cast<TDiscarder*>(Discarder_.Get());
        }

        void Commit() {
            if (!!Sender_ && Y_LIKELY(!!Commiter_)) {
                Sender()->ReleaseStream().Swap(*Commiter());
                Sender_.Destroy();
                Commiter_ = THolder<TThreadedQueue::ICallback>(Queue_->Run(Commiter_.Release(), TInstant::Max()).Release());

                if (Y_UNLIKELY(!Commiter_)) {
                    return;
                }

                if (Commiter()->Exception()) {
                    std::rethrow_exception(Commiter()->Exception());
                }
            }
        }

        void Discard() noexcept {
            if (!!Sender_ && Y_LIKELY(!!Discarder_)) {
                Sender()->ReleaseStream().Swap(*Discarder());
                Sender_.Destroy();
                Discarder_ = THolder<TThreadedQueue::ICallback>(Queue_->Run(Discarder_.Release(), TInstant::Max()).Release());

                if (Y_UNLIKELY(!!Discarder_)) {
                    return;
                }

                if (Commiter()->Exception()) {
                    std::rethrow_exception(Commiter()->Exception());
                }
            }
        }

    private:
        TError DoSend(TChunkList lst, TInstant) noexcept override {
            try {
                try {
                    if (!Sender_) {
                        return {};
                    }

                    if (lst.Empty()) {
                        Commit();
                        return {};
                    }

                    Sender()->Swap(lst);
                    Sender_ = THolder<TThreadedQueue::ICallback>(Queue_->Run(Sender_.Release(), TInstant::Max()).Release());

                    if (Y_UNLIKELY(!Sender_)) {
                        return {};
                    }

                    if (Sender()->Exception()) {
                        std::rethrow_exception(Sender()->Exception());
                    }

                    if (Sender()->Empty()) {
                        Sender()->Cleanup();
                        Sender()->Swap(lst);
                    }
                    return {};
                } catch (...) {
                    Sender_.Destroy();
                    throw;
                }
            } Y_TRY_STORE(yexception);
        }

    private:
        TThreadedQueue* const Queue_{ nullptr };

        THolder<TThreadedQueue::ICallback> Sender_;
        THolder<TThreadedQueue::ICallback> Commiter_;
        THolder<TThreadedQueue::ICallback> Discarder_;
    };

    const THolder<TThreadedQueue::ICallback> callback(
        queue->Run<TThreadedQueue::ICallback>(new TStore(cacher, meta, id), deadline).Release());

    if (!callback) {
        return nullptr;
    }

    TStore* const store = static_cast<TStore*>(callback.Get());

    if (store->Exception()) {
        std::rethrow_exception(store->Exception());
    }

    if (!*store) {
        return nullptr;
    }

    return MakeHolder<TOutput>(cacher, queue, std::move(*store));
}

bool NCache::MakeSpace(TMemoryCacher* cacher, TThreadedQueue* queue, TInstant deadline, const TString& id, size_t size) {
    class TMakeSpace
        : public TThreadedQueue::ICallback
        , public TNonCopyable
    {
    public:
        TMakeSpace(const TMemoryCacher* const cacher, const TString& id, size_t size) noexcept
            : Cacher_(cacher)
            , Id_(id)
            , Size_(size)
            , Result_(false)
        {}

        bool Result() const noexcept {
            return Result_;
        }

    private:
        void DoRun() noexcept override {
            if (Awaited()) {
                Result_ = Cacher_->MakeSpace(Id_, Size_);
            }
        }

    private:
        const TMemoryCacher* const Cacher_{ nullptr };
        const TString Id_;
        const size_t Size_;
        bool Result_;
    };

    const THolder<TThreadedQueue::ICallback> callback(
        queue->Run<TThreadedQueue::ICallback>(new TMakeSpace(cacher, id, size), deadline).Release());

    return !!callback && static_cast<const TMakeSpace*>(callback.Get())->Result();
}
