// Copyright (C) 2024 The Qt Company Ltd.
// SPDX-License-Identifier: LicenseRef-Qt-Commercial OR GPL-3.0-only

#ifndef RPCJOBS_H
#define RPCJOBS_H

#include "qrpcbench_common.h"

#include <grpc/support/time.h>
#include <proto/bench.grpc.pb.h>
#include <proto/bench.pb.h>

#include <grpcpp/server_context.h>

using qt::bench::BenchmarkService;

class RpcJob
{
public:
    RpcJob() = default;
    virtual ~RpcJob() = default;

    RpcJob(const RpcJob &) = delete;
    RpcJob &operator=(const RpcJob &) = delete;

    RpcJob(RpcJob &&) = delete;
    RpcJob &operator=(RpcJob &&) = delete;

    virtual void onProcess(bool ok) = 0;
    virtual void onDone() = 0;

    [[nodiscard]] bool isDone() const noexcept { return mOnDoneCalled; }

private:
    std::atomic<bool> mOnDoneCalled = false;
    friend class NotifyWhenDone;
};

class NotifyWhenDone final : public RpcJob
{
public:
    explicit NotifyWhenDone(RpcJob *parent) : mParent(parent) { }

    void onProcess(bool /* ok */) override
    {
        assert(mParent);
        mParent->mOnDoneCalled = true;
        onDone();
    }

    void onDone() override { delete this; }

private:
    RpcJob *mParent;
};

class UnaryCall final : public RpcJob
{
    using Request = qt::bench::UnaryCallRequest;
    using Response = qt::bench::UnaryCallResponse;

public:
    explicit UnaryCall(grpc::ServerCompletionQueue *cq, BenchmarkService::AsyncService *service)
        : mCQ(cq), mService(service), mWriter(&mContext)
    {
        mService->RequestUnaryCall(&mContext, &mRequest, &mWriter, cq, cq, this);
        mContext.AsyncNotifyWhenDone(new NotifyWhenDone(this));
    }

    void onProcess(bool ok) override
    {
        if (!ok || isDone())
            return onDone();

        if (mState == State::Created) {
            new UnaryCall(mCQ, mService);
            *mResponse.mutable_timestamp() = getTimestamp();
            mResponse.set_request_latency_nanos(calculateLatencyNanos(mRequest.timestamp(),
                                                                      mResponse.timestamp()));
            if (mRequest.has_payload())
                mResponse.mutable_payload()->swap(*mRequest.mutable_payload());
            mWriter.Finish(mResponse, grpc::Status::OK, this);
            mState = State::Processed;
        } else {
            return onDone();
        }
    }

    void onDone() override { delete this; }

private:
    grpc::ServerContext mContext;
    grpc::ServerCompletionQueue *mCQ;
    BenchmarkService::AsyncService *mService;

    Request mRequest;
    Response mResponse;
    grpc::ServerAsyncResponseWriter<Response> mWriter;

    enum class State { Created, Processed };
    State mState = State::Created;
};

class ServerStreaming final : public RpcJob
{
    using Request = qt::bench::ServerStreamingRequest;
    using Response = qt::bench::ServerStreamingResponse;

public:
    explicit ServerStreaming(grpc::ServerCompletionQueue *cq,
                             BenchmarkService::AsyncService *service)
        : mCQ(cq), mService(service), mServerStream(&mContext)
    {
        mService->RequestServerStreaming(&mContext, &mRequest, &mServerStream, mCQ, mCQ, this);
        mContext.AsyncNotifyWhenDone(new NotifyWhenDone(this));
    }

    void onProcess(bool ok) override
    {
        if (!ok || isDone())
            return onDone();

        switch (mState) {
        case State::Created: {
            new ServerStreaming(mCQ, mService);

            if (mRequest.has_payload())
                mResponse.mutable_payload()->assign(mRequest.payload().size(), 'x');
            mResponse.set_pong(mPongCount);

            mServerStream.Write(mResponse, this);
            mState = State::Writing;
        } break;
        case State::Writing: {
            if (++mPongCount >= mRequest.ping()) {
                mServerStream.Finish(grpc::Status::OK, this);
                mState = State::Finished;
            } else {
                mResponse.set_pong(mPongCount);
                mServerStream.Write(mResponse, this);
            }
        } break;
        case State::Finished: {
            // Finish will invoke onDone directly.
        } break;
        }
    }

    void onDone() override
    {
        std::cout << std::format("serverstream done, {} pongs\n", mPongCount);
        delete this;
    }

private:
    grpc::ServerContext mContext;
    grpc::ServerCompletionQueue *mCQ;
    BenchmarkService::AsyncService *mService;
    grpc::ServerAsyncWriter<Response> mServerStream;

    enum class State { Created, Writing, Finished };
    State mState = State::Created;

    Request mRequest;
    Response mResponse;
    uint64_t mPongCount = 0;
};

class ClientStreaming final : public RpcJob
{
    using Request = qt::bench::ClientStreamingRequest;
    using Response = qt::bench::ClientStreamingResponse;

public:
    explicit ClientStreaming(grpc::ServerCompletionQueue *cq,
                             BenchmarkService::AsyncService *service)
        : mCQ(cq), mService(service), mClientStream(&mContext)
    {
        mService->RequestClientStreaming(&mContext, &mClientStream, mCQ, mCQ, this);
        mContext.AsyncNotifyWhenDone(new NotifyWhenDone(this));
    }

    void onProcess(bool ok) override
    {
        if (isDone())
            return onDone();

        switch (mState) {
        case State::Created: {
            new ClientStreaming(mCQ, mService);
            if (!ok)
                return onDone();
            mClientStream.Read(&mRequest, this);
            mState = State::Reading;
        } break;
        case State::Reading: {
            if (!ok) {
                mResponse.set_pong(mPingCount);
                mClientStream.Finish(mResponse, grpc::Status::OK, this);
                mState = State::Finished;
            } else {
                ++mPingCount;
                if (mRequest.has_payload() && !mResponse.has_payload())
                    mResponse.mutable_payload()->assign(mRequest.payload().size(), 'x');
                mRequest.set_ping(mPingCount);
                mClientStream.Read(&mRequest, this);
            }
        } break;
        case State::Finished: {
            // Finish will invoke onDone directly.
        } break;
        }
    }

    void onDone() override
    {
        std::cout << std::format("clientstream done, {} pings\n", mPingCount);
        delete this;
    }

private:
    grpc::ServerContext mContext;
    grpc::ServerCompletionQueue *mCQ;
    BenchmarkService::AsyncService *mService;
    grpc::ServerAsyncReader<Response, Request> mClientStream;

    enum class State { Created, Reading, Finished };
    State mState = State::Created;

    Request mRequest;
    Response mResponse;
    uint64_t mPingCount = 0;
};

class BiDiStreaming final : public RpcJob
{
    using Request = qt::bench::BiDiStreamingRequest;
    using Response = qt::bench::BiDiStreamingResponse;
public:
    explicit BiDiStreaming(grpc::ServerCompletionQueue *cq, BenchmarkService::AsyncService *service)
        : mCQ(cq), mService(service), mStream(&mContext)
    {
        mService->RequestBiDiStreaming(&mContext, &mStream, mCQ, mCQ, this);
        mContext.AsyncNotifyWhenDone(new NotifyWhenDone(this));
    }

    void onProcess(bool ok) override
    {
        if (isDone())
            return onDone();

        switch (mState) {
        case State::Created: {
            new BiDiStreaming(mCQ, mService);
            if (!ok)
                return onDone();
            auto md = mContext.client_metadata();
            if (const auto it = md.find("write-queries"); it == md.cend()) {
                std::cout << "requested 'write-queries' not found in metadata\n";
                return onDone();
            } else {
                mWriteQueries = std::stoul({ it->second.data() });
            }
            if (const auto it = md.find("write-size"); it != md.cend())
                mWriteSize = std::stoul({ it->second.data() });
            mState = State::Processing;
            mReadOperation = new ReadOperation(this);
            mWriteOperation = new WriteOperation(this);
        } break;
        case State::Processing: {
        } break;
        case State::Finished: {
            // Finish will invoke onDone directly.
        } break;
        }
    }

    void onDone() override
    {
        std::cout << "bidistream done\n";
        delete mReadOperation;
        delete mWriteOperation;
        delete this;
    }

private:
    grpc::ServerContext mContext;
    grpc::ServerCompletionQueue *mCQ;
    BenchmarkService::AsyncService *mService;

    grpc::ServerAsyncReaderWriter<Response, Request> mStream;

    enum class State { Created, Processing, Finished };
    State mState = State::Created;

    class ReadOperation final : public RpcJob
    {
    public:
        explicit ReadOperation(BiDiStreaming *parent) : mParent(parent)
        {
            mParent->mStream.Read(&mRequest, this);
        }

        ~ReadOperation() override = default;

        void onProcess(bool ok) override
        {
            if (mParent->mState != State::Processing)
                return;

            if (!ok) {
                mParent->mReadDone = true;
                if (mParent->mWriteDone) { // both ends finished
                    mParent->mState = State::Finished;
                    mParent->mStream.Finish(grpc::Status::OK, mParent);
                }
                return;
            }

            mParent->mStream.Read(&mRequest, this);
        }

        void onDone() override { delete this; }

    private:
        BiDiStreaming *mParent;
        Request mRequest;
    };

    class WriteOperation final : public RpcJob
    {
    public:
        explicit WriteOperation(BiDiStreaming *parent) : mParent(parent)
        {
            if (mParent->mWriteSize > 0)
                mResponse.set_payload(std::string(mParent->mWriteSize, 'x'));
            mParent->mStream.Write(mResponse, this);
            ++mWritten;
        }

        ~WriteOperation() override = default;

        void onProcess(bool ok) override
        {
            if (mParent->mState == State::Finished)
                return;

            if (mWritten >= mParent->mWriteQueries) {
                if (mParent->mReadDone || !ok) { // both ends finished
                    mParent->mState = State::Finished;
                    mParent->mStream.Finish(grpc::Status::OK, mParent);
                }
                return;
            }

            mParent->mStream.Write(mResponse, this);
            ++mWritten;
        }

        void onDone() override { delete this; }

    private:
        BiDiStreaming *mParent;
        Response mResponse;
        size_t mWritten = 0;
    };

    friend class ReadOperation;
    friend class WriteOperation;
    ReadOperation *mReadOperation{};
    WriteOperation *mWriteOperation{};
    size_t mWriteQueries = 0;
    size_t mWriteSize = 0;
    std::atomic<bool> mReadDone = false;
    std::atomic<bool> mWriteDone = false;
};

#endif // RPCJOBS_H
