#include "http2_stream.h"
#include "http2_edge_prio_fix.h"

#include <util/generic/strbuf.h>

namespace NSrvKernel::NHTTP2 {
    namespace {
        void BuildConnDescrProps(TConnProps& props, IConnection& conn, THTTP11Request& req) noexcept {
            const auto& mainProps = *conn.GetMainDescr().Properties;
            props.Start = TInstant::Now();
            props.Random = RandomNumber<ui64>();
            props.SocketIo = nullptr;
            props.UserConnIsSsl = mainProps.UserConnIsSsl && "https" == req.Scheme;
        }
    }


    TStream::TBody::TBody(TStream& stream, IConnection& conn, THTTP11Request request) noexcept
        : Request(std::move(request))
        , ClientInput(stream, Request.Props(), conn)
        , ClientOutput(stream, conn, stream.PrioTreeNode_)
    {
        IdempotentRequest = IsIn({EMethod::HEAD, EMethod::GET}, Request.RequestLine().Method);
    }

    TStream::TStream(IConnection& conn, const ui32 streamId) noexcept
        : TStreamId(streamId)
        , TaskName_("http2stream#" + ::ToString(StreamId_))
        , Conn_(conn)
        , Logger_(Conn_.GetLogger())
        , Stats_(Conn_.GetStats())
        , AuxServerSettings_(conn.GetAuxServerSettings())
        , PrioTreeNode_(Conn_.GetPrioTreeRoot(), streamId)
    {
        Y_HTTP2_METH_E(Logger_);
    }

    TStream::~TStream() {
        Y_HTTP2_METH_E(Logger_);
        Y_VERIFY(!BodyPtr_);
    }

    TError TStream::Open(const TFrameHeading headersHeading, THTTP11Request request, size_t sendBufferMax) noexcept {
        Y_HTTP2_METH(Logger_, PrintReqSpecHeaders(request));

        Stats_.Stream.OnClientOpen();
        Start(Conn_.GetExecutor());

        BodyPtr_ = MakeHolder<TBody>(*this, Conn_, std::move(request));
        BodyPtr_->Request.Props().HTTP2 = &Details_;
        BodyPtr_->SendBufferMax = sendBufferMax;

        if (headersHeading.HasFlagEndStream()) {
            Y_PROPAGATE_ERROR(BodyPtr_->ClientInput.Close());
        }

        return {};
    }

    bool TStream::IsClosedByClient() const noexcept {
        // RFC 7540: half-closed (local):
        //               An endpoint can receive any type of frame in this state.
        //           half-closed (remote):
        //               If an endpoint receives additional frames, other than
        //               WINDOW_UPDATE, PRIORITY, or RST_STREAM, for a stream that is in
        //               this state, it MUST respond with a stream error (Section 5.4.2) of
        //               type STREAM_CLOSED.
        //           closed:
        //               An endpoint that receives any frame other than PRIORITY
        //               after receiving a RST_STREAM MUST treat that as a stream error
        //               (Section 5.4.2) of type STREAM_CLOSED.  Similarly, an endpoint
        //               that receives any frames after receiving a frame with the
        //               END_STREAM flag set MUST treat that as a connection error
        //               (Section 5.4.1) of type STREAM_CLOSED
        //
        // RFC 7540: An endpoint can end a connection at any time.  In particular, an
        //           endpoint MAY choose to treat a stream error as a connection error.
        //
        // Thus we are allowed to be more strict about some errors

        switch (Details_.FinalState) {
        case EFinalState::Unknown:
            return BodyPtr_ && BodyPtr_->ClientInput.IsClosed();
            // RFC 7540 allows us to be more strict about some errors
        case EFinalState::ResetByClient:
        case EFinalState::Closed:
            return true;
        default:
            return false;
        }
    }

    void TStream::PrintTo(IOutputStream& out) const {
        if (BodyPtr_) {
            Y_HTTP2_PRINT_OBJ(out, StreamId_, Running(), Details_.FinalState, PrioTreeNode_,
                BodyPtr_->ClientInput, BodyPtr_->ClientOutput, BodyPtr_->SendBufferMax, BodyPtr_->ParserState);
        } else {
            Y_HTTP2_PRINT_OBJ(out, StreamId_, Running(), Details_.FinalState, PrioTreeNode_,
                (bool)BodyPtr_);
        }
    }

    TError TStream::OnStreamError(std::pair<EErrorCode, TStreamErrorReason> error) noexcept {
        Y_HTTP2_METH(Logger_, error.first, error.second);

        if (IsReset()) {
            Y_HTTP2_EVENT(Logger_, "Already reset", error.first, Details_.FinalState);
        } else {
            Y_HTTP2_BLOCK(Logger_, "Resetting", error.first);
            Details_.FinalState = EFinalState::ResetByServer;
            Details_.ErrorCode = error.first;
            if (!Conn_.IsConnError()) {
                // Being unable to send means the connection is dead
                Y_PROPAGATE_ERROR(
                    Conn_.GetClientOutput().SendRstStream(GetStreamId(), error.first, error.second)
                );
            } else {
                Y_HTTP2_EVENT_E(Logger_, "Skipped sending RST_STREAM");
            }
        }

        // Cancelling any earlier will prevent us from enqueueing our RST_SERVER
        Cancel();

        return {};
    }

    // recving

    TError TStream::DoRecv(TChunkList& lst, TInstant deadline) noexcept {
        Y_HTTP2_METH_E(Logger_);

        Y_REQUIRE(!IsReset(),
                  TSystemError{ECONNRESET});

        return BodyPtr_->ClientInput.DequeueData(deadline).AssignTo(lst);
    }

    TError TStream::OnData(TFrame frame) noexcept {
        Y_HTTP2_METH(Logger_, frame.Heading);

        if (!IsReset()) {
            Y_REQUIRE(IsOpen(),
                TConnectionError(EErrorCode::INTERNAL_ERROR, EConnInternalError::UnexpectedData));
            Y_PROPAGATE_ERROR(BodyPtr_->ClientInput.OnData(std::move(frame)));
        }

        return {};
    }

    TError TStream::OnTrailers(THeadersList) noexcept {
        Y_HTTP2_METH_E(Logger_);

        Y_REQUIRE(!IsClosedByClient(),
            TConnectionError(EErrorCode::STREAM_CLOSED, EConnStreamClosed::ClosedByClient));

        // OnRecvEndStream();
        // TODO (velavokr): implement trailers properly
        return Y_MAKE_ERROR(
            TStreamError(EErrorCode::INTERNAL_ERROR, EStreamInternalError::TrailersNotImplemented));
    }

    TError TStream::OnPriority(TPriority prio, TStream::TPtr streamPtr) noexcept {
        Y_HTTP2_METH(Logger_, prio);

        // RFC 7540: A stream cannot depend on itself.  An endpoint MUST treat this as a
        //           stream error (Section 5.4.2) of type PROTOCOL_ERROR.
        if (streamPtr.Get() == this) {
            if (IsIdle()) {
                // RFC 7540: RST_STREAM frames MUST NOT be sent for a stream in the "idle" state.
                // RFC 7540: endpoint MAY choose to treat a stream error as a connection error.
                return Y_MAKE_ERROR(TConnectionError(EErrorCode::PROTOCOL_ERROR, EConnProtocolError::InvalidPriority));
            } else {
                return Y_MAKE_ERROR(TStreamError(EErrorCode::PROTOCOL_ERROR, EStreamProtocolError::InvalidPriority));
            }
        }

        // The stream ranking may become completely disrupted after the priority change.
        // We must save the contents of the stream output queue to rebuild it later from scratch.
        // We need not handle exceptions thrown between Save and Restore because they immediately cause a GOAWAY.
        Conn_.GetClientOutput().SaveStreamQueue();

        if (streamPtr) {
            Y_HTTP2_EVENT_E(Logger_, "Priority with parent");
            Y_VERIFY(PrioTreeNode_.UpdatePrio(streamPtr->PrioTreeNode_, prio.RawWeight, prio.Exclusive));
        } else if (!prio.StreamDependency) {
            Y_HTTP2_EVENT_E(Logger_, "Priority with root");
            Y_VERIFY(PrioTreeNode_.UpdatePrio(Conn_.GetPrioTreeRoot(), prio.RawWeight, prio.Exclusive));
        } else {
            Y_HTTP2_EVENT_E(Logger_, "Priority with dead parent");
            // RFC 7540: If a stream identified in a dependency has no associated priority
            //           information, then the dependent stream is instead assigned a default
            //           priority (Section 5.3.5).
            Y_VERIFY(PrioTreeNode_.UpdatePrio(Conn_.GetPrioTreeRoot(), RFC_PRIO_RAW_WEIGHT_DEFAULT, false));
        }

        // Rebuilding the stream output queue.
        // See the note above the Save part about why we do not handle exceptions here.
        Conn_.GetClientOutput().RestoreStreamQueue();
        return {};
    }

    void TStream::OnClientRstStream(EErrorCode code) noexcept {
        Y_HTTP2_METH(Logger_, code);

        Cancel();

        // RFC 7540: To avoid looping, an endpoint MUST NOT send a RST_STREAM in response
        //           to a RST_STREAM frame.
        if (IsReset()) {
            Y_HTTP2_EVENT_E(Logger_, "Ignoring in reset state");
        } else {
            Details_.FinalState = EFinalState::ResetByClient;
            Details_.ErrorCode = code;
        }
    }

    void TStream::OnConnError(const TError& connError) noexcept {
        Y_HTTP2_METH_E(Logger_);

        Cancel();

        if (IsReset()) {
            Y_HTTP2_EVENT_E(Logger_, "Ignoring in reset state");
        } else {
            Details_.FinalState = EFinalState::ConnError;

            if (connError) {
                if (const auto* sysErr = connError.GetAs<TSystemError>()) {
                    if (sysErr->Status() == ECANCELED) {
                        Y_HTTP2_EVENT_E(Logger_, "ECANCELED from above");
                        Details_.FinalState = EFinalState::ConnCancel;
                    }
                }
            }
        }
    }

    void TStream::OnReadFin() noexcept {
        Y_HTTP2_METH_E(Logger_);

        if (BodyPtr_) {
            if (BodyPtr_->ClientInput.IsOpen()) {
                Stats_.Stream.OnConnAbort();
                Cancel();
                if (Details_.FinalState == EFinalState::Unknown) {
                    Details_.FinalState = EFinalState::ConnClose;
                }
            }

            BodyPtr_->ClientOutput.GetFlow().OnReadFin();
        }
    }

    // sending

    TError TStream::DoSendHead(TResponse&& response, const bool, TInstant) {
        Y_HTTP2_METH_E(Logger_);

        Y_REQUIRE(!IsReset(),
            TSystemError{EPIPE});

        if (Conn_.GetEdgePrioFix().IsEnabled()) {
            Y_HTTP2_BLOCK_E(Logger_, "Checking response headers");
            if (auto prio = Conn_.GetEdgePrioFix().OnResponse(response.Headers())) {
                Y_HTTP2_BLOCK(Logger_, "Fixing priority in Edge", *this, prio);
                Y_PROPAGATE_ERROR(OnPriority(*prio, nullptr));
            }
        }

        auto& parserState = BodyPtr_->ParserState;
        Y_VERIFY(EParserState::Start == parserState);

        parserState = HTTP_CONTINUE == response.ResponseLine().StatusCode
                      ? EParserState::Headers100Continue
                      : EParserState::Headers;

        const auto contentLength = response.Props().ContentLength;
        const bool eos = EParserState::Headers == parserState &&
                         (0 == contentLength ||
                         EMethod::HEAD == BodyPtr_->Request.RequestLine().Method ||
                         response.ResponseLine().NullContent());

        auto headers = ConvertResponse11To2(Logger_, Stats_, std::move(response));
        Y_HTTP2_BLOCK(Logger_, "Response headers", PrintRespAllHeaders(headers));

        Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.SendHeaders(std::move(headers),
                                                             contentLength, eos));

        if (eos) {
            Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.MarkClosed());
        }

        return {};
    }

    TError TStream::DoSend(TChunkList lst, TInstant) noexcept {
        Y_HTTP2_METH(Logger_, lst.ChunksCount(), lst.size());

        Y_REQUIRE(!IsReset(),
            TSystemError{EPIPE});

        auto& parserState = BodyPtr_->ParserState;
        Y_VERIFY(IsIn({EParserState::Headers, EParserState::Headers100Continue}, parserState));

        if (lst.Empty()) {
            if (EParserState::Headers == parserState) {
                parserState = EParserState::End;

                // We are done sending once we return from the call
                Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.SendEnd());
            } else {
                // BALANCER-1375
                // There is a contract between the balancer modules
                // that every http response must end with an empty chunk list.
                // In case of 100 response we are expecting another response
                // thus we ignore the chunk.
                //
                // RFC 7231: A server that sends a 100 (Continue) response MUST ultimately send
                //           a final status code, once the message body is received and
                //           processed, unless the connection is closed prematurely.
                parserState = EParserState::Start;
            }
        } else  {
            // The call may end before the data is actually sent
            Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.SendData(
                std::move(lst),
                BodyPtr_->SendBufferMax
            ));
        }
        return {};
    }

    TError TStream::DoSendTrailers(THeaders&& headers, TInstant) {
        Y_HTTP2_METH(Logger_, headers.Size());

        Y_REQUIRE(!IsReset(),
                  TSystemError{EPIPE});

        // Converting THeaders to THeadersList
        THeadersList headersList;
        for (const auto &[name, valueList] : headers) {
            for (const auto &value: valueList) {
                headersList.Add(
                        TChunkList{NewChunkForceCopy(name.AsStringBuf())},
                        TChunkList{NewChunkForceCopy(value.AsStringBuf())}
                );
            }
        }

        Y_HTTP2_BLOCK(Logger_, "Trailer headers", PrintRespAllHeaders(headersList));
        Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.FlushData());
        return BodyPtr_->ClientOutput.SendHeaders(std::move(headersList), std::nullopt, true);
    }

    TError TStream::OnInitialWindowUpdate() noexcept {
        Y_HTTP2_METH_E(Logger_);

        if (BodyPtr_) {
            Y_TRY(TError, error) {
                return BodyPtr_->ClientOutput.GetFlow().UpdateInitialSize();
            } Y_CATCH {
                if (const auto* rstStream = error.GetAs<TStreamError>()) {
                    return OnStreamError({rstStream->ErrorCode, rstStream->ErrorReason});
                } else {
                    return error;
                }
            };
        }

        return {};
    }

    TError TStream::OnWindowUpdate(ui32 windowUpdate) noexcept {
        Y_HTTP2_METH(Logger_, windowUpdate);

        if (BodyPtr_) {
            Y_PROPAGATE_ERROR(BodyPtr_->ClientOutput.GetFlow().Update(windowUpdate));
        }
        return {};
    }

    void TStream::Start(TContExecutor& exec) noexcept {
        Y_HTTP2_METH_E(Logger_);

        Y_ASSERT(!StreamTask_.Running());
        StreamTask_ = TCoroutine{TaskName_.c_str(), &exec, [this] {
            Y_HTTP2_BLOCK(Logger_, "OnHTTP", PrintReqSpecHeaders(BodyPtr_->Request));
            Y_TRY(TError, error) {
                Y_REQUIRE(!IsReset(),
                          TSystemError{ECANCELED});

                auto descrProps = Conn_.GetMainDescr().Properties->Copy();
                BuildConnDescrProps(descrProps, Conn_, BodyPtr_->Request);
                TConnDescr connDescr = Conn_.GetMainDescr().Copy(*this, *this, descrProps);
                connDescr.Request = &BodyPtr_->Request;

                TError err = Conn_.OnHTTP(connDescr);
                // If we get TBanAddressError, then the connection should be dropped
                if (err && err.GetAs<TBanAddressError>()) {
                    return err;
                }
                // If output stream is over, handle errors by sending a RST_STREAM with an error code of NO_ERROR
                // Otherwise, propagate the original error
                if (err && !BodyPtr_->ClientOutput.IsClosed()) {
                    return err;
                }

                Y_REQUIRE(!IsReset(),
                          TSystemError{ECANCELED});

                // Preventing stream hanging
                // See incident SEARCHPRODINCIDENTS-3183
                Y_REQUIRE(!BodyPtr_->ClientOutput.IsOpen(),
                    TConnectionError(EErrorCode::INTERNAL_ERROR, EConnInternalError::StreamNotClosed));

                // RFC 7540: A server can
                //           send a complete response prior to the client sending an entire
                //           request if the response does not depend on any portion of the request
                //           that has not been sent and received.  When this is true, a server MAY
                //           request that the client abort transmission of a request without error
                //           by sending a RST_STREAM with an error code of NO_ERROR after sending
                //           a complete response (i.e., a frame with the END_STREAM flag).
                if (err || (EFinalState::Unknown == Details_.FinalState && BodyPtr_->ClientInput.IsOpen())) {
                    return OnStreamError({EErrorCode::NO_ERROR, EStreamNoError::ResponseFinished});
                }

                return {};
            } Y_CATCH {
                Y_HTTP2_CATCH_E(Logger_, *error);
                OnException(std::move(error), EExcState::TopLevel);
            }
            OnFinish();
        }};
    }

    void TStream::Cancel() noexcept {
        Y_HTTP2_METH_E(Logger_);
        StreamTask_.Cancel();
        if (BodyPtr_) {
            BodyPtr_->ClientInput.Cancel();
            BodyPtr_->ClientOutput.Cancel();
        }
    }

    void TStream::Join() noexcept {
        Y_HTTP2_METH_E(Logger_);
        if (StreamTask_.Running()) {
            StreamTask_.Join();
        } else {
            if (!Finished_) {
                OnFinish();
            }
        }
    }

    void TStream::OnFinish() noexcept {
        Finished_ = true;
        if (BodyPtr_ && BodyPtr_->ClientInput.IsClosed() && BodyPtr_->ClientOutput.IsClosed()) {
            Stats_.Stream.OnSuccess();
        }
        BodyPtr_.Reset();
        if (EFinalState::Unknown == Details_.FinalState) {
            Details_.FinalState = EFinalState::Closed;
        }
        Conn_.DisposeStream(*this);
    }

    void TStream::OnException(TError err, EExcState excState) noexcept {
        Y_TRY(TError, error) {
            Y_HTTP2_METH(Logger_, GetErrorMessage(err));

            if (auto* connErr = err.GetAs<TConnectionError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling TConnectionError");
                Conn_.OnConnError(connErr->ErrorCode, connErr->AsStrBuf(), connErr->ErrorReason);
            } else if (auto* streamErr = err.GetAs<TStreamError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling TStreamError");
                Y_PROPAGATE_ERROR(OnStreamError({streamErr->ErrorCode, streamErr->ErrorReason}));
            } else if (auto* forceClose = err.GetAs<TForceStreamClose>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling TForceStreamClose");
                Y_PROPAGATE_ERROR(OnStreamError(GetStreamBackendError()));
            } else if (auto* systemErr = err.GetAs<TSystemError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling TSystemError");
                if (EExcState::Backend == excState && BodyPtr_ && BodyPtr_->ClientOutput.IsOpen()
                    && IsIn<int>({ETIMEDOUT, ECONNRESET, EIO, EPIPE}, systemErr->Status()))
                {
                    Y_HTTP2_BLOCK_E(Logger_, "Handling a backend error");
                    Y_TRY(TError, flushErr) {
                        return BodyPtr_->ClientOutput.FlushData();
                    } Y_CATCH {
                        Y_HTTP2_CATCH_E(Logger_, *flushErr);
                        OnException(std::move(flushErr), EExcState::FinalFlush);
                    }
                }

                if (EExcState::Backend == excState) {
                    Y_PROPAGATE_ERROR(OnStreamError(GetStreamBackendError()));
                } else {
                    Y_PROPAGATE_ERROR(OnStreamError({
                        EErrorCode::INTERNAL_ERROR,
                        EExcState::Backend == excState ?
                        EStreamInternalError::BackendError :
                        EStreamInternalError::ResponseSendError
                    }));
                }
            } else if (auto* httpErr = err.GetAs<THttpError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling NSrvKernel::THttpError");
                // Balancer raises 400 not only if it is a client's fault
                // but also when it cannot parse the server's response.
                if (EExcState::Backend == excState || httpErr->Code() >= 500) {
                    Y_PROPAGATE_ERROR(OnStreamError(GetStreamBackendError()));
                } else {
                    Y_PROPAGATE_ERROR(OnStreamError({EErrorCode::PROTOCOL_ERROR, EStreamProtocolError::Http4xx}));
                }
            } else if (auto* backErr = err.GetAs<TBackendError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling NSrvKernel::TBackendError");

                TError innerError;
                backErr->StoreTo(innerError);
                if (innerError && EExcState::TopLevel == excState) {
                    OnException(std::move(innerError), EExcState::Backend);
                } else {
                    Y_HTTP2_EVENT(Logger_, "Invalid backend error", this, !!innerError, excState);
                    Conn_.OnConnError(EErrorCode::INTERNAL_ERROR, "Invalid backend error",
                        innerError ?
                            EConnInternalError::EmptyBackendError :
                            EConnInternalError::RecursiveBackendError
                    );
                }
            } else if (auto* banErr = err.GetAs<TBanAddressError>()) {
                Y_HTTP2_BLOCK_E(Logger_, "Handling NSrvKernel::TBanAddressError");
                Conn_.OnConnInternalError(std::move(err));
            } else {
                Y_HTTP2_BLOCK(Logger_, "Handling unknown error", this);
                Stats_.Stream.OnUnknownError();
                Y_PROPAGATE_ERROR(OnStreamError({EErrorCode::INTERNAL_ERROR, EStreamInternalError::UnknownError}));
            }
            return {};
        } Y_CATCH {
            Y_HTTP2_CATCH(Logger_, *error, this);
            Conn_.OnConnInternalError(std::move(error));
        }
    }

    std::pair<EErrorCode, TStreamErrorReason> TStream::GetStreamBackendError() const noexcept {
        if (AuxServerSettings_.RefusedStreamOnErrorEnabled
            && BodyPtr_ && EParserState::Start == BodyPtr_->ParserState && BodyPtr_->IdempotentRequest)
        {
            // BALANCER-2398 Making sure the browsers would retry the request
            return {EErrorCode::REFUSED_STREAM, ERefusedStream::BackendError};
        } else {
            return {EErrorCode::INTERNAL_ERROR, EStreamInternalError::BackendError};
        }
    }

}

Y_HTTP2_GEN_PRINT(TStream);
