#include "http2_connection.h"
#include "http2_stream.h"
#include <balancer/kernel/connection_manager_helpers/helpers.h>
#include <balancer/kernel/coro/universal_sleep.h>

using namespace NBalancerClient;
using namespace NSrvKernel;

namespace {
nghttp2_nv MakeHeader(TStringBuf name, TStringBuf value) {
    return {
            .name = reinterpret_cast<ui8*>(const_cast<TStringBuf::pointer>(name.data())),
            .value = reinterpret_cast<ui8 *>(const_cast<TStringBuf::pointer>(value.data())),
            .namelen = name.size(),
            .valuelen = value.size(),
    };
}
}

std::shared_ptr<nghttp2_session> THttp2Connection::CreateSession(THttp2Connection* http2Connection) {
    auto callbacks = CreateCallbacks();
    nghttp2_session *raw;
    if (int res = nghttp2_session_client_new(&raw, callbacks.get(), http2Connection); res != 0) {
        ythrow yexception{} << nghttp2_strerror(res);
    }
    return {raw, [](auto *ptr) {
        nghttp2_session_del(ptr);
    }};
}

std::shared_ptr<nghttp2_session_callbacks> THttp2Connection::CreateCallbacks() {
    nghttp2_session_callbacks *raw;
    if (int res = nghttp2_session_callbacks_new(&raw); res != 0) {
        ythrow yexception{} << nghttp2_strerror(res);
    }

    nghttp2_session_callbacks_set_on_header_callback2(raw, THttp2Connection::OnHeader);
    nghttp2_session_callbacks_set_on_data_chunk_recv_callback(raw, THttp2Connection::OnData);
    nghttp2_session_callbacks_set_on_frame_recv_callback(raw, THttp2Connection::OnFrameRecv);
    nghttp2_session_callbacks_set_on_frame_send_callback(raw, THttp2Connection::OnFrameSent);
    nghttp2_session_callbacks_set_on_stream_close_callback(raw, THttp2Connection::OnStreamClose);
    nghttp2_session_callbacks_set_send_data_callback(raw, THttp2Connection::OnSendData);

    return {raw, [](auto *ptr) {
        nghttp2_session_callbacks_del(ptr);
    }};
}

ssize_t THttp2Connection::Read(nghttp2_session *session, int32_t streamId, ui8* buf, size_t length, ui32* dataFlags, nghttp2_data_source* /*source*/, void* /*user_data*/) {
    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, streamId));

    if (!stream) {
        return NGHTTP2_ERR_TEMPORAL_CALLBACK_FAILURE;
    }

    return stream->CheckNewData(buf, length, dataFlags);
}

int THttp2Connection::OnFrameRecv(nghttp2_session *session, const nghttp2_frame *frame, void* data) {
    auto* connection = static_cast<THttp2Connection*>(data);
    if (frame->hd.type == NGHTTP2_PING) {
        bool ack = frame->hd.flags & NGHTTP2_FLAG_ACK;
        if (ack) {
            connection->PingContext_.PingAck();
        }
    } else {
        if (frame->hd.type == NGHTTP2_SETTINGS) {
            connection->PrefaceReceivedEvent_.BroadCast();
        }

        connection->PingContext_.ResetSuccessivePings();
    }

    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, frame->hd.stream_id));

    if (!stream) {
        return 0;
    }

    return stream->OnFrameRecv(frame);
}

int THttp2Connection::OnFrameSent(nghttp2_session* session, const nghttp2_frame* frame, void*) {
    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, frame->hd.stream_id));

    if (!stream) {
        return 0;
    }

    return stream->OnFrameSent(frame);
}

int THttp2Connection::OnData(nghttp2_session *session, ui8 flags, int32_t streamId, const ui8 *data, size_t len, void* /*userData*/) {
    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, streamId));
    if (!stream) {
        return 0;
    }
    return stream->OnData(flags, data, len);
}

int THttp2Connection::OnStreamClose(nghttp2_session* session, int32_t streamId, ui32 errorCode, void* user_data) {
    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, streamId));

    TError error;
    if (errorCode != NGHTTP2_NO_ERROR) {
        error = Y_MAKE_ERROR(TBackendError(TStringBuilder{} << "stream closed with error \"" << nghttp2_http2_strerror(errorCode) << "\""));
    }

    auto* connection = static_cast<THttp2Connection*>(user_data);
    connection->OnStreamClose(stream, std::move(error));

    return 0;
}

int THttp2Connection::OnHeader(nghttp2_session *session, const nghttp2_frame *frame, nghttp2_rcbuf *nameRcBuf, nghttp2_rcbuf *valueRcBuf, ui8 flags, void* /*userData*/) {
    if (frame->hd.type == NGHTTP2_PUSH_PROMISE) {
        return NGHTTP2_ERR_TEMPORAL_CALLBACK_FAILURE;
    }

    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, frame->hd.stream_id));

    if (!stream) {
        return 0;
    }

    return stream->OnHeader(frame, nameRcBuf, valueRcBuf, flags);
}

TError THttp2Connection::Start(TInstant deadline) {
    Y_PROPAGATE_ERROR(SubmitSettings());
    Y_PROPAGATE_ERROR(Send(deadline));

    Coroutine_ = TCoroutine{ECoroType::Service, "connection coroutine", RunningCont()->Executor(), std::mem_fn(&THttp2Connection::HandleSocket), this};

    int status = PrefaceReceivedEvent_.WaitD(deadline);
    Y_REQUIRE(status == EWAKEDUP, TSystemError{status} << "preface wasn't received ");

    if (!Error_) {
        Pinger_.Shedule(PingContext_);
    }

    return Error_.Clone();
}

THttp2Connection::THttp2Connection(const TString& addr, THolder<NModProxy::TKeepAliveData> connection, TPinger& pinger)
    : Addr_{addr}
    , BackendConnection_{std::move(connection)}
    , Session_{CreateSession(this)}
    , Pinger_{pinger}
    , PingContext_{TPingContext::TOptions {
        .Pinger = pinger,
        .ContExecutor = RunningCont()->Executor(),
        .PingHandler = [this](TInstant deadline) -> TError {
            Y_REQUIRE(nghttp2_submit_ping(Session_.get(), 0, nullptr) == 0, yexception{} << "ping submit failed");
            return Send(deadline);
        },
        .CancelHandler = [this](TError error) {
            Error_ = std::move(error);
            Coroutine_.Cancel();
        }
    }}
    , PrefaceReceivedEvent_{RunningCont()->Executor()}
{
}

NSrvKernel::TErrorOr<i32> THttp2Connection::SubmitRequest(const NSrvKernel::TRequest& request) {
    TString method = ToString(request.RequestLine().Method);
    auto host = request.Headers().GetFirstValue("Host");
    Y_REQUIRE(host, yexception{} << "Host header is required");
    auto path = request.RequestLine().GetURL();
    TVector<nghttp2_nv> hdrs{
            MakeHeader(":method", method),
            MakeHeader(":scheme", BackendConnection_->HaveSsl() ? "https" : "http"),
            MakeHeader(":authority", host),
            MakeHeader(":path", path)};

    for (auto& header: request.Headers()) {
        if (AsciiEqualsIgnoreCase(header.first.AsStringBuf(), "Host")) {
            continue;
        }
        hdrs.emplace_back(MakeHeader(header.first.AsStringBuf(), header.second.front().AsStringBuf()));
    }

    nghttp2_data_provider data;
    data.read_callback = THttp2Connection::Read;
    i32 streamId = nghttp2_submit_request(Session_.get(),
                                          nullptr,
                                          hdrs.data(),
                                          hdrs.size(),
                                          request.Props().ContentLength == 0 ? nullptr : &data,
                                          nullptr);
    Y_REQUIRE(streamId >= 0, yexception{} << nghttp2_strerror(streamId));
    return streamId;
}

NSrvKernel::TError THttp2Connection::SubmitSettings() {
    static const ui32 window_size = 1 << 20;

    std::array<nghttp2_settings_entry, 2> settings;
    settings[0].settings_id = NGHTTP2_SETTINGS_MAX_FRAME_SIZE;
    settings[0].value = 1 << 22;

    settings[1].settings_id = NGHTTP2_SETTINGS_INITIAL_WINDOW_SIZE;
    settings[1].value = window_size;

    {
        int res = nghttp2_submit_settings(Session_.get(), NGHTTP2_FLAG_NONE, settings.data(), settings.size());
        Y_REQUIRE(res == 0, yexception{} << nghttp2_strerror(res));
    }

    {
        int res = nghttp2_session_set_local_window_size(Session_.get(), NGHTTP2_FLAG_NONE, 0, window_size);
        Y_REQUIRE(res == 0, yexception{} << nghttp2_strerror(res));
    }

    return {};
}

NSrvKernel::TErrorOr<THolder<THttp2Stream>> THttp2Connection::StartStream(const TRequest& request,
                                                                          IIoInput& input,
                                                                          IHttpOutput& output,
                                                                          THttp2Stream::TCallbacks callbacks,
                                                                          TInstant deadline) {
    Y_REQUIRE(Alive(), TBackendError{"connection is not alive"});
    i32 streamId;
    Y_PROPAGATE_ERROR(SubmitRequest(request).AssignTo(streamId));
    THolder stream{new THttp2Stream(*this,
                                       streamId,
                                       RunningCont(),
                                       input,
                                       output,
                                       std::move(callbacks))};
    Y_VERIFY(nghttp2_session_set_stream_user_data(Session_.get(), streamId, stream.Get()) == 0, "cannot set stream user data");
    Streams_.PushBack(stream.Get());
    Y_PROPAGATE_ERROR(Send(deadline));

    return stream;
}

void THttp2Connection::HandleSocket() {
    TError error = [this]() -> TError {
        while (nghttp2_session_want_read(Session_.get())) {
            TChunkList chunkList;
            Y_PROPAGATE_ERROR(NModProxy::WrapIntoBackendError(BackendConnection_->Input().Recv(chunkList, TInstant::Max())));
            if (chunkList.Empty()) {
                return Y_MAKE_ERROR(TBackendError{"socket closed"});
            }

            PingContext_.DataReceived();

            while (!chunkList.Empty()) {
                auto chunk = chunkList.PopFront();
                TStringBuf left = chunk->AsStringBuf();
                while (!left.empty()) {
                    ssize_t size = nghttp2_session_mem_recv(Session_.get(),
                                                            reinterpret_cast<const ui8*>(left.data()),
                                                            left.size());
                    if (size < 0) {
                        return Y_MAKE_ERROR(yexception{} << nghttp2_strerror(size));
                    }
                    left = left.substr(size);
                }
            }

            Y_PROPAGATE_ERROR(Send(TInstant::Max()));
        }

        return {};
    }();

    if (IsCancelledError(error) && Error_) {
        error = Error_.Clone();
    }

    if (!Streams_.Empty() && !error) {
        error = Y_MAKE_ERROR(TBackendError{"unknown http2 error"});
    }

    while (!Streams_.Empty()) {
        OnStreamClose(Streams_.PopFront(), error.Clone());
    }

    PrefaceReceivedEvent_.BroadCast();
}

NSrvKernel::TError THttp2Connection::CancelStream(THttp2Stream& stream, TInstant deadline) {
    OnStreamClose(&stream, Y_MAKE_ERROR(yexception{} << "stream cancelled"));
    Y_REQUIRE(nghttp2_submit_rst_stream(Session_.get(), NGHTTP2_FLAG_NONE, stream.GetStreamId(), NGHTTP2_CANCEL) == 0, yexception{} << "failed to reset stream");
    Y_PROPAGATE_ERROR(Send(deadline));
    return {};
}

NSrvKernel::TError THttp2Connection::StreamHasNewData(THttp2Stream& stream, bool resume) {
    if (resume) {
        Y_REQUIRE(nghttp2_session_resume_data(Session_.get(), stream.GetStreamId()) == 0, yexception{} << "failed to resume stream data");
    }
    Y_PROPAGATE_ERROR(Send(TInstant::Max()));
    return {};
}

int THttp2Connection::OnSendData(nghttp2_session* session, nghttp2_frame* frame, const ui8* framehd, size_t length, nghttp2_data_source* source, void* user_data) {
    Y_UNUSED(source);

    auto* connection = static_cast<THttp2Connection*>(user_data);
    if (connection->DataSent_) {
        connection->DataSent_ = false;
        return 0;
    }

    auto* stream = static_cast<THttp2Stream*>(nghttp2_session_get_stream_user_data(session, frame->hd.stream_id));
    if (!stream) {
        return NGHTTP2_ERR_TEMPORAL_CALLBACK_FAILURE;
    }

    TChunkList data;
    if (auto error = stream->GetNewData(length).AssignTo(data)) {
        return NGHTTP2_ERR_TEMPORAL_CALLBACK_FAILURE;
    }

    connection->DataToSent_.PushNonOwning(TStringBuf{reinterpret_cast<const char*>(framehd), 9});
    if (frame->data.padlen > 0) {
        ui8 padlen = frame->data.padlen - 1;
        TString buf{reinterpret_cast<char*>(&padlen), 1};
        connection->DataToSent_.Push(std::move(buf));
    }
    connection->DataToSent_.Append(std::move(data));
    if (frame->data.padlen > 1) {
        TString buf{frame->data.padlen - 1, 0};
        connection->DataToSent_.Push(std::move(buf));
    }

    return NGHTTP2_ERR_WOULDBLOCK;
}

NSrvKernel::TError THttp2Connection::Send(TInstant deadline) {
    if (SendInProgress_) {
        return {};
    }
    SendInProgress_ = true;
    Y_DEFER {
        SendInProgress_ = false;
    };
    while (nghttp2_session_want_write(Session_.get())) {
        char* data;
        ssize_t size = nghttp2_session_mem_send(Session_.get(),
                                                const_cast<const ui8**>(reinterpret_cast<ui8**>(&data)));
        if (size < 0) {
            return Y_MAKE_ERROR(yexception{} << nghttp2_strerror(size));
        }

        if (size == 0 && DataToSent_.Empty()) {
            break;
        }

        TChunkList chunkList;
        if (DataToSent_.Empty()) {
            chunkList.PushNonOwning(TStringBuf{data, static_cast<size_t>(size)});
        } else {
            chunkList.Swap(DataToSent_);
            DataSent_ = true;
        }
        Y_PROPAGATE_ERROR(NModProxy::WrapIntoBackendError(BackendConnection_->Output().Send(std::move(chunkList), deadline)));
    }

    return {};
}

TString THttp2Connection::ResolvedAddr() const {
    return BackendConnection_->AddrStr();
}

bool THttp2Connection::Alive() const  {
    return Coroutine_.Running();
}

bool THttp2Connection::HasStreams() const {
    return !Streams_.Empty();
}

TMaybe<TDuration> THttp2Connection::IdleTime() const {
    if (HasStreams()) {
        return Nothing();
    }

    return Now() - LastStreamTime_;
}

const TString& THttp2Connection::Addr() const {
    return Addr_;
}

bool THttp2Connection::HasStream(i32 streamId) {
    return nghttp2_session_get_stream_user_data(Session_.get(), streamId) != nullptr;
}

void THttp2Connection::OnStreamClose(THttp2Stream* stream, TError error) {
    if (stream) {
        Y_VERIFY(nghttp2_session_set_stream_user_data(Session_.get(), stream->GetStreamId(), nullptr) == 0, "cannot set stream user data");
        stream->OnStreamClose(std::move(error));
    }

    if (!HasStreams()) {
        LastStreamTime_ = Now();
    }
}
