#include "http2_stream_out.h"
#include "http2_edge_prio_fix.h"
#include "http2_stats.h"
#include "http2_stream.h"

#include <balancer/kernel/memory/split.h>

namespace NSrvKernel::NHTTP2 {

    TStreamOutput::TStreamOutput(TStreamId streamId, IConnection& conn, const TPrioTreeNode& prioTreeNode) noexcept
        : IStreamOutput(streamId.GetStreamId())
        , Logger_(conn.GetLogger())
        , Stats_(conn.GetStats())
        , Output_(conn.GetClientOutput())
        , PrioTreeNode_(prioTreeNode)
        , PrioHandle_(*this)
        , Executor_(conn.GetExecutor())
        , Flow_(conn.GetLogger(), *this, conn.GetClientSettings())
    {}

    void TStreamOutput::PrintTo(IOutputStream& out) const {
        Y_HTTP2_PRINT_OBJ(
            out, StreamId_, State_, Flow_, Output_.GetFlow(), DataBuffer_.size(), PrioHandle_.IsLinked()
        );
    }

    TError TStreamOutput::SendHeaders(THeadersList headers, std::optional<ui64> contentLength, bool eos) noexcept {
        Y_HTTP2_METH(Logger_, PrintRespSpecHeaders(headers));

        Y_REQUIRE(EState::Open == State_,
            TSystemError{EPIPE});

        ContentLength_ = contentLength;

        return Output_.SendHeaders(GetStreamId(), std::move(headers), eos);
    }

    TError TStreamOutput::SendData(TChunkList data, size_t sendBufferMax) noexcept {
        Y_HTTP2_METH(Logger_, data.ChunksCount(), data.size(), sendBufferMax);

        Y_REQUIRE(MaySendData(),
            TSystemError{EPIPE});

        if (ContentLength_) {
            auto len = data.size();
            Y_REQUIRE(len <= *ContentLength_,
                THttpError{HTTP_BAD_REQUEST});
            *ContentLength_ -= len;

            if (!*ContentLength_) {
                State_ = EState::Closing;
            }
        }

        // ToDO(dudkomv): optimize copying
        data.MakeOwned();
        DataBuffer_.Append(std::move(data));

        if (!DataBuffer_.Empty()) {
            do {
                Y_HTTP2_BLOCK(Logger_, "Requesting async data send", this);

                while (Flow_.IsBlocked() && MaySendData()) {
                    const int ret = SendDataCV_.wait(&Executor_);
                    Y_REQUIRE(ret == EWAKEDUP,
                              TSystemError{ret});
                }

                // The TConnOutput could have emptied DataBuffer and maybe even marked the stream as closed by now.
                if (DataBuffer_.Empty()) {
                    break;
                }

                Y_REQUIRE(MaySendData(),
                    TSystemError{EPIPE});

                Y_PROPAGATE_ERROR(Output_.SendDataAsync(PrioHandle_));

                Y_HTTP2_BLOCK(Logger_, "Waiting for async data send", this);

                while (DataBuffer_.size() > sendBufferMax && PrioHandle_.IsLinked() && MaySendData()) {
                    const int ret = SendDataCV_.wait(&Executor_);
                    Y_REQUIRE(ret == EWAKEDUP,
                              TSystemError{ret});
                }
            } while (DataBuffer_.size() > sendBufferMax);

            // One final check before return
            if (!DataBuffer_.Empty()) {
                Y_REQUIRE(MaySendData(),
                    TSystemError{EPIPE});

                if (!Flow_.IsBlocked()) {
                    Y_PROPAGATE_ERROR(Output_.SendDataAsync(PrioHandle_));
                }
            }
        }
        return {};
    }

    TError TStreamOutput::FlushData() noexcept {
        Y_HTTP2_METH_E(Logger_);
        Y_PROPAGATE_ERROR(SendData(TChunkList(), 0));
        PrioHandle_.Unlink();
        return {};
    }

    TError TStreamOutput::MarkClosed() noexcept {
        Y_HTTP2_METH_E(Logger_);

        if (EState::Closed == State_) {
            return {};
        }

        Y_REQUIRE(MaySendData(),
            TSystemError{EPIPE});

        State_ = EState::Closed;
        Stats_.Stream.OnServerClose();

        return {};
    }

    TError TStreamOutput::SendEnd() noexcept {
        Y_HTTP2_METH_E(Logger_);

        if (EState::Closed == State_) {
            return {};
        }

        Y_REQUIRE(MaySendData(),
            TSystemError{EPIPE});

        State_ = EState::Closing;
        Y_PROPAGATE_ERROR(FlushData());
        // By now the PrioHandle is guaranteed to be unlinked

        if (EState::Closing == State_) {
            // Only us can send EOS by now
            Y_PROPAGATE_ERROR(Output_.SendDataEndStream(GetStreamId()));
            Stats_.Stream.OnServerClose();
            State_ = EState::Closed;
        }

        return {};
    }

    void TStreamOutput::Cancel() noexcept {
        Y_HTTP2_METH_E(Logger_);
        if (State_ != EState::Closed) {
            State_ = EState::Cancelled;
            SendDataCV_.notify();
            PrioHandle_.Unlink();
        }
    }

    TErrorOr<TData> TStreamOutput::DequeueNextDataFrame(ui32 frameSize) noexcept {
        Y_HTTP2_METH_E(Logger_);

        if (!MaySendData()) {
            Y_HTTP2_EVENT_E(Logger_, "Cannot send");
            return TData();
        }

        TData data;

        data.Data = CutPrefix(
            Min(Flow_.GetAvailableSize(), frameSize),
            DataBuffer_
        );

        if (DataBuffer_.Empty() && EState::Closing == State_) {
            Y_HTTP2_EVENT_E(Logger_, "Sending EOS with the last frame");
            data.EndOfStream = true;
            Stats_.Stream.OnServerClose();
            State_ = EState::Closed;
        }

        Y_PROPAGATE_ERROR(Flow_.Consume(data.Data.size()));
        SendDataCV_.notify();

        Y_HTTP2_EVENT(Logger_, "Dequeued", data);

        return data;
    }

    void TStreamOutput::OnConnFlowBlockedForever() noexcept {
        Y_HTTP2_METH_E(Logger_);
        SendDataCV_.notify();
    }

    void TStreamOutput::OnConnOutputClose() noexcept {
        Y_HTTP2_METH_E(Logger_);
        SendDataCV_.notify();
    }

    void TStreamOutput::OnFlowBlockedForever() noexcept {
        Y_HTTP2_METH_E(Logger_);
        SendDataCV_.notify();
    }

    TError TStreamOutput::OnFlowUnblocked() noexcept {
        Y_HTTP2_METH_E(Logger_);
        SendDataCV_.notify();

        // BALANCER-1364
        if (MaySendData() && !DataBuffer_.Empty()) {
            // We are relying on the invariant:
            // if the stream is neither closed nor cancelled and while DataBuffer has some data
            // it is perfectly safe to call SendDataAsync once again.
            // SendDataAsync cannot throw since the conditions on which it may
            // are mutually exclusive with the event we are handling now.
            Y_PROPAGATE_ERROR(Output_.SendDataAsync(PrioHandle_));
        }
        return {};
    }

    bool TStreamOutput::MaySendData() const noexcept {
        return IsIn<EState>({EState::Open, EState::Closing}, State_)
            && Output_.IsOpen()
            && !Flow_.IsBlockedForever()
            && !Output_.GetFlow().IsBlockedForever();
    }
}

Y_HTTP2_GEN_PRINT(TStreamOutput);
