// Copyright (c) 2026 Roman Ilin
// SPDX-License-Identifier: BUSL-1.1
//
// Use of this software is governed by the Business Source License
// included in the LICENSE.txt file.

#include <arpa/inet.h>
#include <errno.h>
#include <inttypes.h>
#include <linux/net_tstamp.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdbit.h>
#include <stddef.h>
#include <stddefer.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/eventfd.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <time.h>
#include <unistd.h>

#include <liburing.h>
#include <sodium.h>

#if __STDC_ENDIAN_NATIVE__ == __STDC_ENDIAN_BIG__
#error "Big-endian hardware is not supported"
// #define MESSAGE_NORMALIZE(v) ({ \
//     static_assert( \
//         sizeof(v) == 2 || \
//         sizeof(v) == 4 || \
//         sizeof(v) == 8 || \
//         sizeof(v) == 16, \
//         "MESSAGE_NORMALIZE: Unsupported value size" \
//     ); \
//     __builtin_choose_expr(sizeof(v) == 2, __builtin_bswap16(v), \
//     __builtin_choose_expr(sizeof(v) == 4, __builtin_bswap32(v), \
//     __builtin_choose_expr(sizeof(v) == 8, __builtin_bswap64(v), \
//     __builtin_choose_expr(sizeof(v) == 16, __builtin_bswap128(v), \
//     __builtin_unreachable())))); \
// })
// #else
// #define MESSAGE_NORMALIZE(v) (v)
#endif

typedef enum : uint8_t {
    MESSAGE_PING = 0,
    MESSAGE_PONG = 1,
} Message_Type;

typedef uint64_t Correlation_ID;

typedef struct {
    Correlation_ID correlation_id;
    Message_Type type;
    uint8_t _pad[7];
} Message_Header;

typedef uint64_t Node_ID;

#define ID_FMT "016" PRIx64

typedef struct {
    int64_t s;
    int32_t ns;
    uint8_t _pad[4];
} Timestamp;

static constexpr int32_t NS_PER_S = 1'000'000'000;

static inline Timestamp timestamp_normalize(int64_t s, int64_t ns)
{
    Timestamp t;
    int64_t s_adjust = ns / NS_PER_S;
    ns = ns % NS_PER_S;
    if (ns < 0) {
        s_adjust--;
        ns += NS_PER_S;
    }
    t.s = s + s_adjust;
    t.ns = (int32_t)ns;
    return t;
}

static Timestamp timestamp_from_timespec(const struct timespec *ts_p)
{
    return timestamp_normalize((int64_t)ts_p->tv_sec, (int64_t)ts_p->tv_nsec);
}

static Timestamp timestamp_now(clockid_t clock_id) {
    struct timespec ts;
    if (clock_gettime(clock_id, &ts) != 0) return (Timestamp){};
    return timestamp_from_timespec(&ts);
}

typedef struct {
    Node_ID node_id;
} Message_Ping;

typedef struct {
    Node_ID node_id;
    Timestamp timestamp_rx;
    Timestamp timestamp_tx;
} Message_Pong;

typedef struct {
    Message_Header header;
    union {
        Message_Ping ping;
        Message_Pong pong;
    } payload;
} Message;

static size_t packet_handle(
    Node_ID node_id,
    int thread_id,
    const struct sockaddr_storage *client_address_p,
    socklen_t address_length,
    uint8_t *buffer_p,
    size_t rx_length,
    size_t max_buffer_length,
    Timestamp rx_timestamp
) 
{
    (void)thread_id;
    (void)client_address_p;
    (void)address_length;
    (void)max_buffer_length;

    if (rx_length < sizeof(Message_Header)) return 0;

    auto message_p = (Message *)buffer_p;

    switch (message_p->header.type) {
        case MESSAGE_PING: {
            if (rx_length != sizeof(Message_Header) + sizeof(Message_Ping)) {
                fprintf(stderr, "Invalid PING message length %zu\n", rx_length);
                return 0;
            }

            Node_ID tx_node_id = message_p->payload.ping.node_id;
            printf("Received PING from node %" ID_FMT "\n", tx_node_id);

            message_p->header.type = MESSAGE_PONG;
            message_p->payload.pong.node_id = node_id;
            message_p->payload.pong.timestamp_rx = rx_timestamp;
            message_p->payload.pong.timestamp_tx = timestamp_now(CLOCK_REALTIME);

            return sizeof(Message_Header) + sizeof(Message_Pong);
        }

        case MESSAGE_PONG: {
            Timestamp tx_timestamp = message_p->payload.pong.timestamp_tx;
            printf("Message type: %d, TX timestamp: %zu.%d\n", message_p->header.type, tx_timestamp.s, tx_timestamp.ns);
            return 0;
        }

        default:
            printf("Unknown message type: %d\n", message_p->header.type);
            return 0;
    }
}

static constexpr uint16_t PORT = 12421;
static constexpr int BGID = 1;
static constexpr size_t BUFFERS_LENGTH = 1024;
static constexpr size_t MTU = 1200;
static constexpr size_t BUFFER_LENGTH = 2048;
static constexpr uint32_t SQ_THREAD_IDLE_MS = 2000;

static int shutdown_efd_g = -1;
static atomic_bool keep_running_g = true;

typedef enum {
    QE_RX,
    QE_SEND,
    QE_BUFFER,
    QE_SHUTDOWN,
} QE_Type;

typedef struct {
    QE_Type type;
    int bid;
} QE_Ctx;

typedef union {
    struct io_uring_recvmsg_out hdr;
    uint8_t raw[BUFFER_LENGTH];
} Aligned_Buffer;

typedef struct {
    Node_ID node_id;
    int thread_id;
    int sfd;
    struct io_uring ring;
    Aligned_Buffer buffers[BUFFERS_LENGTH];
    QE_Ctx ctxs[BUFFERS_LENGTH];
    struct msghdr send_messages[BUFFERS_LENGTH];
    struct iovec send_iovs[BUFFERS_LENGTH];
    struct msghdr rx_message_template;
    QE_Ctx rx_ctx;
    QE_Ctx shutdown_ctx;
} Worker_Ctx;

static void signal_handle(int sig)
{
    (void)sig;
    keep_running_g = false;
    if (shutdown_efd_g >= 0) {
        uint64_t val = 1;
        ssize_t res = write(shutdown_efd_g, &val, sizeof(val));
        (void)res;
    }
}

static struct io_uring_sqe *sqe_get_safe(struct io_uring *ring_p)
{
    auto sqe_p = io_uring_get_sqe(ring_p);
    if (!sqe_p) {
        fprintf(stderr, "Fatal: SQ ring is full\n");
        abort();
    }
    return sqe_p;
}

static void buffer_provide(Worker_Ctx *worker_p, int bid)
{
    auto sqe_p = sqe_get_safe(&worker_p->ring);
    worker_p->ctxs[bid].type = QE_BUFFER;
    worker_p->ctxs[bid].bid = bid;
    io_uring_prep_provide_buffers(sqe_p, worker_p->buffers[bid].raw, BUFFER_LENGTH, 1, BGID, bid);
    io_uring_sqe_set_data(sqe_p, &worker_p->ctxs[bid]);
}

static void *worker_thread(void *arg_p)
{
    Worker_Ctx *worker_p = calloc(1, sizeof(Worker_Ctx));
    defer free(worker_p);

    worker_p->node_id = 0; // TODO
    worker_p->thread_id = *(int *)arg_p;
    worker_p->rx_ctx.type = QE_RX;
    worker_p->rx_ctx.bid = -1;
    worker_p->shutdown_ctx.type = QE_SHUTDOWN;
    worker_p->shutdown_ctx.bid = -1;
    worker_p->rx_message_template.msg_namelen = sizeof(struct sockaddr_storage);

    worker_p->sfd = socket(AF_INET, SOCK_DGRAM, 0);
    defer if (worker_p->sfd >= 0) close(worker_p->sfd);
    if (worker_p->sfd < 0) {
        perror("Failed to create socket");
        return nullptr;
    }

    auto socket_option_reuse_port = 1;
    if (setsockopt(
        worker_p->sfd,
        SOL_SOCKET,
        SO_REUSEPORT,
        &socket_option_reuse_port,
        sizeof(socket_option_reuse_port)
    ) < 0) {
        perror("Failed to set SO_REUSEPORT");
        return nullptr;
    }

    auto socket_option_timestamping =
        SOF_TIMESTAMPING_TX_SOFTWARE |
        SOF_TIMESTAMPING_SOFTWARE |
        SOF_TIMESTAMPING_OPT_ID;
    if (setsockopt(
        worker_p->sfd,
        SOL_SOCKET,
        SO_TIMESTAMPING,
        &socket_option_timestamping,
        sizeof(socket_option_timestamping)
    ) < 0) {
        perror("Failed to set SO_TIMESTAMPING");
        return nullptr;
    }

    struct sockaddr_in address = {
        .sin_family = AF_INET,
        .sin_port = htons(PORT),
        .sin_addr.s_addr = INADDR_ANY
    };
    if (bind(worker_p->sfd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("Failed to bind socket");
        return nullptr;
    }

    struct io_uring_params ring_params = {0};
    ring_params.flags = IORING_SETUP_SQPOLL;
    ring_params.sq_thread_idle = SQ_THREAD_IDLE_MS;
    
    defer io_uring_queue_exit(&worker_p->ring);
    if (io_uring_queue_init_params(BUFFERS_LENGTH * 2, &worker_p->ring, &ring_params) < 0) {
        perror("Failed to setup I/O queues");
        return nullptr;
    }

    auto sqe_p = sqe_get_safe(&worker_p->ring);
    io_uring_prep_provide_buffers(sqe_p, worker_p->buffers, BUFFER_LENGTH, BUFFERS_LENGTH, BGID, 0);
    io_uring_sqe_set_data(sqe_p, nullptr);
    io_uring_submit(&worker_p->ring);

    struct io_uring_cqe *cqe_p;
    io_uring_wait_cqe(&worker_p->ring, &cqe_p);
    io_uring_cqe_seen(&worker_p->ring, cqe_p);

    sqe_p = sqe_get_safe(&worker_p->ring);
    io_uring_prep_recvmsg_multishot(sqe_p, worker_p->sfd, &worker_p->rx_message_template, 0);
    sqe_p->flags |= IOSQE_BUFFER_SELECT;
    sqe_p->buf_group = BGID;
    io_uring_sqe_set_data(sqe_p, &worker_p->rx_ctx);

    sqe_p = sqe_get_safe(&worker_p->ring);
    io_uring_prep_poll_add(sqe_p, shutdown_efd_g, POLLIN);
    io_uring_sqe_set_data(sqe_p, &worker_p->shutdown_ctx);

    io_uring_submit(&worker_p->ring);

    while (keep_running_g) {
        int ret = io_uring_wait_cqe(&worker_p->ring, &cqe_p);
        if (ret < 0) {
            if (ret == -EINTR) continue;
            perror("io_uring_wait_cqe error");
            break;
        }

        unsigned head;
        unsigned count = 0;

        io_uring_for_each_cqe(&worker_p->ring, head, cqe_p) {
            QE_Ctx *ctx_p = io_uring_cqe_get_data(cqe_p);
            
            if (ctx_p && ctx_p->type == QE_RX) {
                if (cqe_p->flags & IORING_CQE_F_BUFFER) {
                    auto bid = (int)(cqe_p->flags >> IORING_CQE_BUFFER_SHIFT);
                    
                    if (cqe_p->res > 0) {
                        auto hdr_p = (struct io_uring_recvmsg_out *)worker_p->buffers[bid].raw;
                        auto client_address_p = io_uring_recvmsg_name(hdr_p);
                        auto payload_p = (uint8_t *)io_uring_recvmsg_payload(hdr_p, &worker_p->rx_message_template);
                        
                        auto actual_payload_length = io_uring_recvmsg_payload_length(
                            hdr_p, cqe_p->res, &worker_p->rx_message_template
                        );

                        struct timespec rx_timespec = {0};
                        struct cmsghdr *cmsg_p = io_uring_recvmsg_cmsg_firsthdr(hdr_p, &worker_p->rx_message_template);
                        while (cmsg_p != nullptr) {
                            if (cmsg_p->cmsg_level == SOL_SOCKET && cmsg_p->cmsg_type == SCM_TIMESTAMPNS) {
                                memcpy(&rx_timespec, CMSG_DATA(cmsg_p), sizeof(struct timespec));
                                break;
                            }
                            cmsg_p = io_uring_recvmsg_cmsg_nexthdr(hdr_p, &worker_p->rx_message_template, cmsg_p);
                        }
                        Timestamp rx_timestamp = timestamp_from_timespec(&rx_timespec);

                        size_t max_payload_length = BUFFER_LENGTH - (size_t)(payload_p - worker_p->buffers[bid].raw);
                        size_t response_length = packet_handle(
                            worker_p->node_id,
                            worker_p->thread_id,
                            (struct sockaddr_storage *)client_address_p,
                            hdr_p->namelen,
                            payload_p,
                            actual_payload_length,
                            max_payload_length,
                            rx_timestamp
                        );

                        if (response_length > 0) {
                            worker_p->send_iovs[bid].iov_base = payload_p;
                            worker_p->send_iovs[bid].iov_len = response_length;

                            worker_p->send_messages[bid].msg_name = client_address_p;
                            worker_p->send_messages[bid].msg_namelen = hdr_p->namelen;
                            worker_p->send_messages[bid].msg_iov = &worker_p->send_iovs[bid];
                            worker_p->send_messages[bid].msg_iovlen = 1;

                            worker_p->ctxs[bid].type = QE_SEND;
                            worker_p->ctxs[bid].bid = bid;

                            auto send_sqe_p = sqe_get_safe(&worker_p->ring);
                            io_uring_prep_sendmsg(send_sqe_p, worker_p->sfd, &worker_p->send_messages[bid], 0);
                            io_uring_sqe_set_data(send_sqe_p, &worker_p->ctxs[bid]);
                        } else {
                            buffer_provide(worker_p, bid);
                        }
                    } else {
                        buffer_provide(worker_p, bid);
                    }
                }

                if (!(cqe_p->flags & IORING_CQE_F_MORE)) {
                    if (cqe_p->res < 0 && cqe_p->res != -ENOBUFS) {
                        usleep(50000);
                    }
                    if (keep_running_g) {
                        auto rearm_sqe_p = sqe_get_safe(&worker_p->ring);
                        io_uring_prep_recvmsg_multishot(rearm_sqe_p, worker_p->sfd, &worker_p->rx_message_template, 0);
                        rearm_sqe_p->flags |= IOSQE_BUFFER_SELECT;
                        rearm_sqe_p->buf_group = BGID;
                        io_uring_sqe_set_data(rearm_sqe_p, &worker_p->rx_ctx);
                    }
                }
            } else if (ctx_p && ctx_p->type == QE_SEND) {
                buffer_provide(worker_p, ctx_p->bid);
            }

            count++;
        }

        io_uring_cq_advance(&worker_p->ring, count);
        io_uring_submit(&worker_p->ring);
    }
    
    return nullptr;
}

int main()
{
    // E.g. for IDs generation
    if (sodium_init() < 0) {
        fprintf(stderr, "Failed to init libsodium\n");
        return EXIT_FAILURE;
    }

    shutdown_efd_g = eventfd(0, EFD_NONBLOCK);
    if (shutdown_efd_g < 0) {
        perror("Failed to create eventfd");
        return EXIT_FAILURE;
    }
    defer close(shutdown_efd_g);

    signal(SIGINT, signal_handle);
    signal(SIGTERM, signal_handle);

    auto cores_count = sysconf(_SC_NPROCESSORS_ONLN);
    pthread_t threads[cores_count];
    int thread_ids[cores_count];

    for (int i = 0; i < cores_count; i++) {
        thread_ids[i] = i;
        pthread_create(&threads[i], nullptr, worker_thread, &thread_ids[i]);
    }

    for (int i = 0; i < cores_count; i++) {
        pthread_join(threads[i], nullptr);
    }

    return EXIT_SUCCESS;
}