// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#ifndef BRPC_RDMA_ENDPOINT_H
#define BRPC_RDMA_ENDPOINT_H

#if BRPC_WITH_RDMA
#include <cstring>
#include <iostream>
#include <string>
#include <vector>
#include <functional>
#include <infiniband/verbs.h>
#include "butil/atomicops.h"
#include "butil/iobuf.h"
#include "butil/macros.h"
#include "butil/containers/mpsc_queue.h"
#include "brpc/socket.h"


namespace brpc {
class Socket;
namespace rdma {

DECLARE_bool(rdma_use_polling);
DECLARE_int32(rdma_poller_num);
DECLARE_bool(rdma_disable_bthread);

class RdmaHandshakeClientV2;
class RdmaHandshakeServerV2;
class RdmaHandshakeClientV3;
class RdmaHandshakeServerV3;
struct ParsedHello;
class RdmaHello;
class RdmaEndpoint;
namespace v2_wire {
    int ReadBodyAndNegotiate(RdmaEndpoint* ep, ParsedHello* remote, bool* negotiated);
    int DrainBytes(RdmaEndpoint* ep, size_t n);
}  // namespace v2_wire

namespace v3_wire {
    void FillLocalRdmaHello(const RdmaEndpoint* ep, RdmaHello* msg);
    int  ReadAndParseV3Hello(RdmaEndpoint* ep, RdmaHello* out);
    int  WriteV3Hello(RdmaEndpoint* ep, const RdmaHello& msg);
}  // namespace v3_wire

class RdmaConnect : public AppConnect {
public:
    void StartConnect(const Socket* socket, 
            void (*done)(int err, void* data), void* data) override;
    void StopConnect(Socket*) override;
    struct RunGuard {
        RunGuard(RdmaConnect* rc) { this_rc = rc; }
        ~RunGuard() { if (this_rc) this_rc->Run(); }
        RdmaConnect* this_rc;
    };

private:
    void Run();
    void (*_done)(int, void*){NULL};
    void* _data{NULL};
};

struct RdmaResource {
    RdmaResource* next{NULL};
    ibv_qp* qp{NULL};
    // For polling mode.
    ibv_cq* polling_cq{NULL};
    // For event mode.
    ibv_cq* send_cq{NULL};
    ibv_cq* recv_cq{NULL};
    ibv_comp_channel* comp_channel{NULL};
    RdmaResource() = default;
    ~RdmaResource();
    DISALLOW_COPY_AND_ASSIGN(RdmaResource);
};

class BAIDU_CACHELINE_ALIGNMENT RdmaEndpoint : public SocketUser {
friend class RdmaConnect;
friend class Socket;
friend class RdmaHandshakeClientV2;
friend class RdmaHandshakeServerV2;
friend class RdmaHandshakeClientV3;
friend class RdmaHandshakeServerV3;
friend int v2_wire::ReadBodyAndNegotiate(RdmaEndpoint*, ParsedHello*, bool*);
friend int v2_wire::DrainBytes(RdmaEndpoint*, size_t);
friend void v3_wire::FillLocalRdmaHello(const RdmaEndpoint*, RdmaHello*);
friend int  v3_wire::ReadAndParseV3Hello(RdmaEndpoint*, RdmaHello*);
friend int  v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&);
public:
    explicit RdmaEndpoint(Socket* s);
    ~RdmaEndpoint() override;

    // Global initialization
    // Return 0 if success, -1 if failed and errno set
    static int GlobalInitialize();

    static void GlobalRelease();

    // Reset the endpoint (for next use)
    void Reset();

    // Cut data from the given IOBuf list and use RDMA to send
    // Return bytes cut if success, -1 if failed and errno set
    ssize_t CutFromIOBufList(butil::IOBuf** data, size_t ndata);

    // Whether the endpoint can send more data
    bool IsWritable() const;

    // For debug
    void DebugInfo(std::ostream& os,
                   butil::StringPiece connector = "\n") const;

    // Callback when there is new epollin event on TCP fd
    static void OnNewDataFromTcp(Socket* m);

    // Initialize polling mode
    static int PollingModeInitialize(bthread_tag_t tag,
                                     std::function<void(void)> callback,
                                     std::function<void(void)> init_fn,
                                     std::function<void(void)> release_fn);

    static void PollingModeRelease(bthread_tag_t tag);

private:
    enum State {
        UNINIT = 0x0,
        C_ALLOC_QPCQ = 0x1,
        C_HELLO_SEND = 0x2,
        C_HELLO_WAIT = 0x3,
        C_BRINGUP_QP = 0x4,
        C_ACK_SEND = 0x5,
        S_HELLO_WAIT = 0x11,
        S_ALLOC_QPCQ = 0x12,
        S_BRINGUP_QP = 0x13,
        S_HELLO_SEND = 0x14,
        S_ACK_WAIT = 0x15,
        ESTABLISHED = 0x100,
        FALLBACK_TCP = 0x200,
        FAILED = 0x300
    };

    // Process handshake at the client
    static void* ProcessHandshakeAtClient(void* arg);

    // Process handshake at the server
    static void* ProcessHandshakeAtServer(void* arg);

    // Allocate resources
    // Return 0 if success, -1 if failed and errno set
    int AllocateResources();

    // Release resources
    void DeallocateResources();

    // Send Imm data to the remote side
    // Arguments:
    //     imm: imm data in the WR
    // Return:
    //     0:   success
    //     -1:  failed, errno set
    int SendImm(uint32_t imm);

    // Try to send pure ACK to the remote side
    // Arguments:
    //     num: the number of rq entry received
    // Return:
    //     0:   success
    //     -1:  failed, errno set
    int SendAck(int num);

    // Handle CQE
    // If wc is not RDMA RECV event:
    //     return 0 if success, -1 if failed and errno set
    // If wc is RDMA RECV event:
    //     return bytes appended if success, -1 if failed and errno set
    ssize_t HandleCompletion(ibv_wc& wc);

    // Post a given number of WRs to Recv Queue
    // If zerocopy is true, reallocate block.
    // Return 0 if success, -1 if failed and errno set
    int PostRecv(uint32_t num, bool zerocopy);

    // Post a WR pointing to the block to the local Recv Queue
    // Arguments:
    //     block: the addr to receive data (ibv_sge.addr)
    //     block_size: the maximum length can be received (ibv_sge.length)
    // Return:
    //     0:   success
    //     -1:  failed, errno set
    int DoPostRecv(void* block, size_t block_size);

    // Read at most len bytes from fd in _socket to data
    // wait for _read_butex if encounter EAGAIN
    // return -1 if encounter other errno (including EOF)
    int ReadFromFd(void* data, size_t len);
    int ReadFromFd(butil::IOPortal* data, size_t len);


    // Write at most len bytes from data to fd in _socket
    // wait for _epollout_butex if encounter EAGAIN
    // return -1 if encounter other errno
    int WriteToFd(void* data, size_t len);

    // Write data to fd in _socket.
    // wait for _epollout_butex if encounter EAGAIN.
    // return -1 if encounter other errno.
    int WriteToFd(butil::IOBuf* data);

    // Copy negotiated remote parameters into the endpoint and compute
    // the SQ/RQ window capacities. Called by both
    // ProcessHandshakeAtClient and ProcessHandshakeAtServer after the
    // peer's hello has been validated.
    void ApplyRemoteHello(const ParsedHello& remote);

    // Bringup the QP from RESET state to RTS state
    // Arguments:
    //     lid: remote LID
    //     gid: remote GID
    //     qp_num: remote QP number
    // Return:
    //     0:   success
    //     -1:  failed, errno set
    int BringUpQp(uint16_t lid, ibv_gid gid, uint32_t qp_num);

    // Get event from comp channel and ack the events
    int GetAndAckEvents(SocketUniquePtr& s);

    // Request completion notification on a send/recv CQ.
    int ReqNotifyCq(bool send_cq);

    // Poll CQ and get the work completion
    static void PollCq(Socket* m);

    // Get the description of current handshake state
    std::string GetStateStr() const;

    // Try to read data on TCP fd in _socket
    void TryReadOnTcp();

    // Add cq socket id to poller
    void PollerAddCqSid();

    // Remove cq socket id to poller
    void PollerRemoveCqSid();

    // Not owner
    Socket* _socket;

    // State of Handshake
    butil::atomic<State> _state;

    // Wire-level handshake protocol version (set by dispatch in
    // ProcessHandshakeAtClient/Server). Aligned with the protocol code:
    //   0 = unnegotiated
    //   2 = v2 "RDMA"
    //   3 = v3 "RDM3"
    int _handshake_version;

    // rdma resource
    RdmaResource* _resource;

    // The number of events requiring ack.
    unsigned int _send_cq_events;
    unsigned int _recv_cq_events;

    // The SocketId which wrap the comp channel of CQ.
    SocketId _cq_sid;

    // Capacity of local Send Queue and local Recv Queue
    uint16_t _sq_size;
    uint16_t _rq_size;

    // Act as sendbuf and recvbuf, but requires no memcpy
    std::vector<butil::IOBuf> _sbuf;
    std::vector<butil::IOBuf> _rbuf;
    // Data address of _rbuf
    std::vector<void*> _rbuf_data;
    // Remote block size for receiving
    uint32_t _remote_recv_block_size;

    // The number of new recv WRs acked to the remote side
    uint16_t _accumulated_ack;
    // The number of WRs sent without solicited flag
    uint16_t _unsolicited;
    // The bytes sent without solicited flag
    uint32_t _unsolicited_bytes;
    // The current index should be used for sending
    uint16_t _sq_current;
    // The number of send WRs not signaled
    uint16_t _sq_unsignaled;
    // The just completed send WR's index
    uint16_t _sq_sent;
    // The just completed recv WR's index
    uint16_t _rq_received;
    // The capacity of local window: min(local SQ, remote RQ)
    uint16_t _local_window_capacity;
    // The capacity of remote window: min(local RQ, remote SQ)
    uint16_t _remote_window_capacity;
    // The number of IMM WRs we can post to the local Send Queue.
    uint16_t _sq_imm_window_size;
    // The number of WRs we can send to remote side.
    butil::atomic<uint16_t> _remote_rq_window_size;
    // The number of WRs we can post to the local Send Queue
    butil::atomic<uint16_t> _sq_window_size;
    // The number of new WRs posted in the local Recv Queue
    butil::atomic<uint16_t> _new_rq_wrs;

    // butex for inform read events on TCP fd during handshake
    butil::atomic<int> *_read_butex;

    DISALLOW_COPY_AND_ASSIGN(RdmaEndpoint);

    // Cq socket id operation type
    struct CqSidOp {
        enum OpType {
            ADD,
            REMOVE,
        };
        SocketId sid;
        OpType type;
    };
    // Poller instance
    struct BAIDU_CACHELINE_ALIGNMENT Poller {
        bthread_t tid{INVALID_BTHREAD};
        butil::MPSCQueue<CqSidOp, butil::ObjectPoolAllocator<CqSidOp>> op_queue;
        // Callback used for io_uring/spdk etc
        std::function<void()> callback;
        // Init and Destroy function
        std::function<void()> init_fn;
        std::function<void()> release_fn;
    };
    // Poller group
    struct BAIDU_CACHELINE_ALIGNMENT PollerGroup {
        PollerGroup() : pollers(FLAGS_rdma_poller_num), running(false) {}
        std::vector<Poller> pollers;
        std::atomic<bool> running;
    };
    static std::vector<PollerGroup> _poller_groups;
};

}  // namespace rdma
}  // namespace brpc

#else  // if BRPC_WITH_RDMA

class RdmaEndpoint { };

#endif  // ifdef USE_RD<A

#endif // BRPC_RDMA_ENDPOINT_H
