#pragma once

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

#include <library/cpp/coroutine/engine/impl.h>
#include <balancer/kernel/coro/poller_event.h>

#include <util/network/socket.h>

class TSocketHolder;

namespace NSrvKernel {
    class TSock {
    public:
        TSock(TSocketHolder* s, TContExecutor* e) noexcept
            : S_(s)
            , E_(e)
        {}

        TSocketHolder* Socket() const noexcept {
            return S_;
        }

        TContExecutor* Executor() const noexcept {
            return E_;
        }

    private:
        TSocketHolder* const S_ = nullptr;
        TContExecutor* const E_ = nullptr;
    };

    enum EPollMode {
        PM_LEVEL_TRIGGERED, /* level_triggered */
        PM_EDGE_TRIGGERED, /* edge_triggered */
    };

    class TSocketIn : public IIoInput {
    public:
        TSocketIn(TSock& sock, EPollMode poolMode) noexcept;
        ~TSocketIn();

        TError WaitRdhup() noexcept;

        void SetOnReadyCallback(std::function<void(int)> callback) noexcept;

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override;

        TChunkPtr NextChunk() noexcept {
            if (!L_) {
                return NewChunkReserve();
            }

            return std::move(L_);
        }

    private:
        TSock& Sock_;

        const bool EdgeTriggered_ = false;

        struct TEdgeTriggeredEvent;
        THolder<TEdgeTriggeredEvent> ReadyRead_;

        TChunkPtr L_;

        std::function<void(int)> ReadyCallback_;
        TPollerEvent ReadyEvent_;
    };


    class TSocketOut :
        public IIoOutput,
        public NCoro::IPollEvent,
        public IUserEvent
    {
    public:
        static constexpr size_t DefaultMaxInQueue = 32 * 1024;
        static constexpr size_t FlushAll = 0;

        TSocketOut(TSock& sock, size_t maxInQueue, EPollMode pollMode) noexcept;

        ~TSocketOut() override;

        TError Flush(size_t keep, TInstant deadline) noexcept;

        size_t BufferSize() const noexcept {
            return W_.size();
        }

        size_t MaxInQueue() const noexcept {
            return MaxInQueue_;
        }

    private:
        TError DoSend(TChunkList lst, TInstant deadline) noexcept override;

        void OnPollEvent(int status) noexcept override;
        void Execute() override {
            OnPollEvent(0);
        }
    private:
        const bool EdgeTriggered_ = false;

        TSock& Sock_;
        TChunkList W_;
        TCont* Wake_ = nullptr;
        size_t MaxInQueue_ = 0;
        int Err_ = 0;
        bool InPoller_ = false;
    };


    class TSocketIo {
    public:
        TSocketIo(TSocketHolder* s, TContExecutor* e, size_t maxInQueue, EPollMode pollMode) noexcept
            : Sock_(s, e)
            , In_(Sock_, pollMode)
            , Out_(Sock_, maxInQueue, pollMode)
        {
        }

        TSocketIn& In() noexcept {
            return In_;
        }

        TSocketOut& Out() noexcept {
            return Out_;
        }

        TSock& Sock() noexcept {
            return Sock_;
        }

    private:
        TSock Sock_;
        TSocketIn In_;
        TSocketOut Out_;
    };
}
