#include "http2_conn_out.h"
#include "http2_headers.h"
#include "http2_settings.h"

#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/http2/server/utils/http2_prio_tree.h>
#include <balancer/kernel/memory/split.h>
#include <balancer/kernel/net/socket.h>

#include <util/generic/strbuf.h>

namespace NSrvKernel::NHTTP2 {

    namespace {
        struct TNginxPrioLess {
            bool operator()(const IStreamOutput& a, const IStreamOutput& b) const {
                const TPrioTreeNode& aPrio = a.GetPrioTreeNode();
                const TPrioTreeNode& bPrio = b.GetPrioTreeNode();
                return aPrio.GetDepth() > bPrio.GetDepth()
                       || aPrio.GetDepth() == bPrio.GetDepth() && aPrio.GetRelWeight() < bPrio.GetRelWeight();
            }
        };
    }

    TConnOutput::TConnOutput(
        IConnection& conn,
        IIoOutput& clientOut,
        TContExecutor& executor,
        TStats& stats,
        TLogger& logger,
        const TClientSettings& clientSettings,
        const TAuxServerSettings& auxServerSettings
    ) noexcept
        : Conn_(conn)
        , ClientOut_(clientOut)
        , Executor_(executor)
        , Stats_(stats)
        , Logger_(logger)
        , ClientSettings_(clientSettings)
        , AuxServerSettings_(auxServerSettings)
        , Encoder_(
            GetHPackEncoderSettings(ClientSettings_, AuxServerSettings_)
        )
        , Flow_(logger, *this, ClientSettings_)
        , StreamQueue_(
            CreateStreamQueue(AuxServerSettings_.StreamsPrioQueueType)
        )
    {}

    void TConnOutput::PrintTo(IOutputStream& out) const {
        Y_HTTP2_PRINT_OBJ(out, State_, FrameBuffer_.size(), StreamQueue_->Size(), Flow_);
    }

    bool TConnOutput::IsOpen() const noexcept {
        return EState::Open == State_;
    }

    void TConnOutput::Close() noexcept {
        Y_HTTP2_METH(Logger_, FrameBuffer_.ChunksCount(), FrameBuffer_.size()/*, BusyBufferCV_.QueueSize()*/);
        State_ = EState::Closed;
        BusyBufferCV_.notify_all();
        EmptyBufferCV_.notify();

        while (!StreamQueue_->Empty()) {
            auto* item = StreamQueue_->GetNext();
            item->Get().OnConnOutputClose();
            StreamQueue_->Remove(*item);
        }
    }

    TError TConnOutput::SendHeaders(ui32 streamId, THeadersList headers, bool eos) noexcept {
        Y_HTTP2_METH(Logger_, streamId);
        const auto buffer = EBuffer::Normal;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));

        // !!Attention!!
        // To avoid racing with SettingsAck
        //      we MUST NOT yield between serializing headers and placing ready frames into the main buffer.

        Y_HTTP2_BLOCK(Logger_, "Encoding headers", streamId, PrintRespSpecHeaders(headers));
        auto headerBlock = Encoder_.Encode(headers);
        {
            auto headersFrame = CutPrefix(GetNextFrameSize(), headerBlock);
            DoSerializeFrame(
                TFrameHeading::NewHeaders(streamId).SetFlagEndHeaders(headerBlock.Empty()).SetFlagEndStream(eos),
                std::move(headersFrame),
                buffer
            );
        }
        TChunkList continuationFrame;
        while (!(continuationFrame = CutPrefix(GetNextFrameSize(), headerBlock)).Empty()) {
            DoSerializeFrame(
                TFrameHeading::NewContinuation(streamId).SetFlagEndHeaders(headerBlock.Empty()),
                std::move(continuationFrame),
                buffer
            );
        }
        return {};
    }

    TError TConnOutput::SendDataAsync(IStreamOutput::TPrioHandle& streamHandle) noexcept {
        Y_HTTP2_METH(Logger_, streamHandle.Get().GetStreamId());

        Y_REQUIRE(IsOpen() && !Flow_.IsBlockedForever(),
            TSystemError{EPIPE});

        if (!streamHandle.IsLinked()) {
            StreamQueue_->Insert(streamHandle);
        }

        EmptyBufferCV_.notify();
        return {};
    }

    TError TConnOutput::SendDataEndStream(ui32 streamId) noexcept {
        Y_HTTP2_METH(Logger_, streamId);
        const auto buffer = EBuffer::Normal;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));

        DoSerializeFrame(
            TFrameHeading::NewDataEndStream(streamId),
            TChunkList(),
            buffer
        );
        return {};
    }

    TError TConnOutput::SendStreamWindowUpdate(ui32 streamId, ui32 windowUpdate) noexcept {
        Y_HTTP2_METH(Logger_, streamId, windowUpdate);
        const auto buffer = EBuffer::OOBCtrl;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));

        DoSerializeFrame(
            TFrameHeading::NewStreamWindowUpdate(streamId),
            WriteWindowUpdate(windowUpdate),
            buffer
        );
        return {};
    }

    TError TConnOutput::SendRstStream(ui32 streamId, EErrorCode errorCode, TStreamErrorReason reason) noexcept {
        Y_HTTP2_METH(Logger_, streamId, errorCode);
        const auto buffer = EBuffer::Ctrl;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));

        Stats_.RstStreamSend.OnError(errorCode, reason);

        DoSerializeFrame(
            TFrameHeading::NewRstStream(streamId),
            WriteRstStream(errorCode),
            buffer
        );
        return {};
    }

    void TConnOutput::SendSettings(TServerSettings& serverSettings) noexcept {
        Y_HTTP2_METH_E(Logger_);
        const auto buffer = EBuffer::Ctrl;

        Y_VERIFY(serverSettings.HasPending());

        DoSerializeFrame(
            TFrameHeading::NewSettings(),
            serverSettings.WritePending(),
            buffer
        );
    }

    TError TConnOutput::SendSettingsAck() noexcept {
        Y_HTTP2_METH_E(Logger_);
        const auto buffer = EBuffer::Ctrl;
        // !!Attention!!
        // To avoid racing with SettingsAck
        //      the sending methods MUST NOT yield between forming frames
        //      (e.g. last checking flows, serializing, encoding headers)
        //      and putting them into the main buffer.
        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));

        DoSerializeFrame(
            TFrameHeading::NewSettingsAck(),
            TChunkList(),
            buffer
        );
        return {};
    }

    void TConnOutput::SendPing(TPing ping) noexcept {
        Y_HTTP2_METH(Logger_, ping);

        DoSerializeFrame(
            TFrameHeading::NewPing(),
            ping.Write(),
            EBuffer::OOBCtrl
        );
    }

    TError TConnOutput::SendPingAck(TChunkPtr ping) noexcept {
        Y_HTTP2_METH_E(Logger_);
        const auto buffer = EBuffer::OOBCtrl;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));
        DoSerializeFrame(
            TFrameHeading::NewPingAck(),
            std::move(ping),
            buffer
        );

        // OOBCtrl prepends frames instead of appending them, thus the inverted ordering
        if (AuxServerSettings_.ReciprocalPingEnabled) {
            DoSerializeFrame(
                TFrameHeading::NewPing(),
                TPing(EPingType::PongRTT, TInstant::Now()).Write(),
                buffer
            );
        }
        return {};
    }

    void TConnOutput::SendGoAwayGraceful(TGoAway goAway, TConnErrorReason reason) noexcept {
        Y_HTTP2_METH_E(Logger_);
        DoSendGoAway(std::move(goAway), std::move(reason), EBuffer::Ctrl);
    }

    void TConnOutput::SendGoAwayUrgent(TGoAway goAway, TConnErrorReason reason) noexcept {
        Y_HTTP2_METH_E(Logger_);
        FrameBuffer_.Clear();
        DoSendGoAway(std::move(goAway), std::move(reason), EBuffer::OOBCtrl);
        Close();
    }

    void TConnOutput::SendShutdownWarning() noexcept {
        Y_HTTP2_METH_E(Logger_);
        // OOBCtrl prepends frames instead of appending them, thus the inverted ordering

        TPing ping{EPingType::GuardGoAwayNoId, TInstant::Now()};
        Y_HTTP2_EVENT(Logger_, "Prepending PING", ping);
        DoSerializeFrame(TFrameHeading::NewPing(), ping.Write(), EBuffer::OOBCtrl);

        TGoAway goaway{RFC_STREAM_ID_MAX, EErrorCode::NO_ERROR};
        Y_HTTP2_EVENT(Logger_, "Prepending GOAWAY", goaway);
        DoSerializeFrame(TFrameHeading::NewGoAway(), goaway.Write(), EBuffer::OOBCtrl);
    }

    TError TConnOutput::SendConnWindowUpdate(ui32 windowUpdate) noexcept {
        Y_HTTP2_METH(Logger_, windowUpdate);
        const auto buffer = EBuffer::OOBCtrl;

        Y_PROPAGATE_ERROR(DoWaitBuffer(buffer));
        DoSerializeFrame(
            TFrameHeading::NewConnectionWindowUpdate(),
            WriteWindowUpdate(windowUpdate),
            buffer
        );
        return {};
    }

    void TConnOutput::SaveStreamQueue() noexcept {
        Y_HTTP2_METH_E(Logger_);
        Y_VERIFY(StreamQueueBackup_.Empty());
        StreamQueueBackup_.MoveFrom(*StreamQueue_);
    }

    void TConnOutput::RestoreStreamQueue() noexcept {
        Y_HTTP2_METH_E(Logger_);
        Y_VERIFY(StreamQueue_->Empty());
        StreamQueue_->MoveFrom(StreamQueueBackup_);
    }

    void TConnOutput::DoSendGoAway(TGoAway goAway, TConnErrorReason reason, TConnOutput::EBuffer buffer) noexcept {
        Y_HTTP2_METH(Logger_, goAway.ErrorCode, reason, buffer);

        auto lst = TChunkList(goAway.Write());

        if (AuxServerSettings_.GoAwayDebugDataEnabled) {
            auto debugData = CutPrefix(
                Max<size_t>(lst.size(), GetNextCtrlFrameSize()) - lst.size(),
                goAway.DebugData
            );
            lst.Append(std::move(debugData));
        }

        Stats_.GoAwaySend.OnError(goAway, reason);

        DoSerializeFrame(
            TFrameHeading::NewGoAway(),
            std::move(lst),
            buffer
        );
    }

    TError TConnOutput::DoWaitBuffer(TConnOutput::EBuffer buffer) noexcept {
        Y_HTTP2_METH_E(Logger_);

        const size_t bufferMax = AuxServerSettings_.ConnDataSendBufferMax
                                + (EBuffer::Normal == buffer ? 0 : AuxServerSettings_.ConnCtrlSendBufferMax);

        while (IsOpen() && FrameBuffer_.size() >= bufferMax) {
            const int ret = BusyBufferCV_.wait(&Executor_);
            Y_REQUIRE(ret == EWAKEDUP,
                      TSystemError{ret});
        }

        // TODO (velavokr): do we need this check here?
        Y_REQUIRE(IsOpen(),
                  TSystemError{EPIPE});
        return {};
    }

    void TConnOutput::DoSerializeFrame(TFrameHeading heading, TChunkPtr payload, TConnOutput::EBuffer buffer) noexcept {
        Y_HTTP2_METH_E(Logger_);
        DoSerializeFrame(heading, TChunkList(std::move(payload)), buffer);
    }

    void TConnOutput::DoSerializeFrame(TFrameHeading heading, TChunkList payload, TConnOutput::EBuffer buffer) noexcept {
        Y_HTTP2_METH(Logger_, heading, payload.size(), buffer);

        heading.Length = payload.size();
        payload.PushFront(heading.Write());

        if (EBuffer::OOBCtrl == buffer) {
            FrameBuffer_.Prepend(std::move(payload));
        } else {
            FrameBuffer_.Append(std::move(payload));
        }

        // Signalling is cheap. So no point in making it conditional.
        EmptyBufferCV_.notify();
    }

    ui32 TConnOutput::GetBufferFreeSpace() const noexcept {
        const ui32 length = FrameBuffer_.size();
        return std::max(length, AuxServerSettings_.ConnDataSendBufferMax) - length;
    }

    ui32 TConnOutput::GetNextCtrlFrameSize() const noexcept {
        return std::min(
            ClientSettings_.MaxFrameSize,
            AuxServerSettings_.ServerFrameSizeMax
        );
    }

    ui32 TConnOutput::GetNextFrameSize() const noexcept {
        return std::min(
            GetNextCtrlFrameSize(),
            GetBufferFreeSpace()
        );
    }

    ui32 TConnOutput::GetNextDataFrameSize() const noexcept {
        return std::min(
            std::min(
                GetNextFrameSize(),
                AuxServerSettings_.ServerDataFrameSizeMax
            ),
            Flow_.GetAvailableSize()
        );
    }

    void TConnOutput::OnFlowBlockedForever() noexcept {
        Y_HTTP2_METH_E(Logger_);
        EmptyBufferCV_.notify();
    }

    TError TConnOutput::OnFlowUnblocked() noexcept {
        Y_HTTP2_METH_E(Logger_);
        EmptyBufferCV_.notify();
        return {};
    }

    void TConnOutput::Start() noexcept {
        Y_HTTP2_METH_E(Logger_);

        Y_ASSERT(!ConnTask_.Running());
        ConnTask_ = TCoroutine{ECoroType::Service, "http2output", &Executor_, [this] {
            Y_TRY(TError, error) {
                // While there is something to send in the future or just now.
                while (IsOpen() || !FrameBuffer_.Empty() || !StreamQueue_->Empty()) {
                    while (IsOpen() && FrameBuffer_.Empty() &&
                           (StreamQueue_->Empty() || !Flow_.IsBlockedForever() && !GetNextDataFrameSize())) {
                        const int ret = EmptyBufferCV_.wait(&Executor_);
                        Y_REQUIRE(ret == EWAKEDUP,
                                  TSystemError{ret});
                    }

                    while (!StreamQueue_->Empty() && (Flow_.IsBlockedForever() || GetNextDataFrameSize())) {
                        Y_HTTP2_BLOCK(Logger_, "Processing stream priority queue", Flow_, GetNextDataFrameSize());
                        auto* item = StreamQueue_->GetNext();

                        // !!Attention!!
                        // To avoid racing with SettingsAck
                        //      we MUST NOT yield between checking conn and stream flows and placing ready frames into the main queue.

                        TData frame;
                        Y_PROPAGATE_ERROR(
                            item->Get().DequeueNextDataFrame(GetNextDataFrameSize()).AssignTo(frame));
                        if (!frame.Data.Empty()) {
                            Y_HTTP2_BLOCK(Logger_, "Trying to send next data frame", item->Get(), frame);

                            if (Flow_.IsBlockedForever()) {
                                Y_HTTP2_BLOCK(Logger_, "Send flow is blocked forever", Flow_);
                                // The stream will wake and cancel itself.
                                item->Get().OnConnFlowBlockedForever();
                                StreamQueue_->Remove(*item);
                            } else {
                                Y_HTTP2_BLOCK(Logger_, "Sending next data", frame);
                                Y_PROPAGATE_ERROR(Flow_.Consume(frame.Data.size()));
                                DoSerializeFrame(
                                    TFrameHeading::NewData(
                                        item->Get().GetStreamId()
                                    ).SetFlagEndStream(
                                        frame.EndOfStream
                                    ),
                                    std::move(frame.Data),
                                    EBuffer::Normal
                                );

                                StreamQueue_->Update(*item);
                            }
                        } else {
                            StreamQueue_->Remove(*item);
                        }
                    }

                    if (!FrameBuffer_.Empty()) {
                        if (AuxServerSettings_.FlushPingEnabled
                            && (FrameBuffer_.size() + Conn_.GetSocketIo()->Out().BufferSize())
                                > AuxServerSettings_.SockSendBufferSizeMax)
                        {
                            Y_HTTP2_BLOCK(Logger_, "Measuring flush RTT",
                                          FrameBuffer_.size(), Conn_.GetSocketIo()->Out().BufferSize());

                            DoSerializeFrame(
                                TFrameHeading::NewPing(),
                                TPing(EPingType::FlushRTT, TInstant::Now()).Write(),
                                EBuffer::OOBCtrl
                            );
                        }

                        Y_HTTP2_BLOCK(Logger_, "Flushing buffer", FrameBuffer_.size());
                        // Otherwise an underlying stream and another coroutine may race both modifying the FrameBuffer.
                        Y_PROPAGATE_ERROR(ClientOut_.Send(std::move(FrameBuffer_), TInstant::Max()));
                        Y_PROPAGATE_ERROR(Conn_.GetSocketIo()->Out().Flush(AuxServerSettings_.SockSendBufferSizeMax, TInstant::Max()));
                        BusyBufferCV_.notify_all();
                    }
                }

                if (Conn_.GetSocketIo()) {
                    Y_HTTP2_BLOCK_E(Logger_, "Flushing socket");
                    Y_PROPAGATE_ERROR(Conn_.GetSocketIo()->Out().Flush(TSocketOut::FlushAll, TInstant::Max()));
                }
                return {};
            } Y_CATCH {
                Y_HTTP2_CATCH(Logger_, *error, this);
                Conn_.OnConnReset(std::move(error), EErrorSource::SendFrames);
            }
            Close();
        }};
    }

    void TConnOutput::Cancel() noexcept {
        ConnTask_.Cancel();
    }

    void TConnOutput::Join() noexcept {
        if (ConnTask_.Running()) {
            ConnTask_.Join();
        } else {
            Close();
        }
    }

    TConnOutput::TStreamQueuePtr TConnOutput::CreateStreamQueue(EStreamsPrioQueueType qt) noexcept {
        switch (qt) {
        case EStreamsPrioQueueType::FIFO:
            return MakeHolder<TStaticFIFOQueue<IStreamOutput>>();
        case EStreamsPrioQueueType::RoundRobin:
            return MakeHolder<TRoundRobinQueue<IStreamOutput>>();
        case EStreamsPrioQueueType::Nginx:
            return MakeHolder<THeapQueueBase<IStreamOutput, TNginxPrioLess>>();
        }
    }
}

Y_HTTP2_GEN_PRINT(TConnOutput);
