// receiver_hwts_logger.c
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <inttypes.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <net/if.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <linux/net_tstamp.h>
#include <linux/sockios.h>
#include <linux/errqueue.h>
#include <time.h>
#include <sys/stat.h>
#include <sched.h>
#include <sys/mman.h>

#define RECV_PORT 5000
#define MAXLINE   2048

/* --- broadcaster header (must match sender) --- */
#define MAGIC_RBRC 0x52425243u /* 'RBRC' */
#define VERSION    1u

static inline uint64_t ntohll(uint64_t x) {
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
    return ((uint64_t)ntohl(x & 0xFFFFFFFFULL) << 32) | ntohl((uint32_t)(x >> 32));
#else
    return x;
#endif
}

#pragma pack(push, 1)
typedef struct {
    uint32_t magic;        /* MAGIC_RBRC (network order) */
    uint32_t version;      /* VERSION    (network order) */
    uint64_t seq_be;       /* sequence number (network order) */
    uint64_t send_sec_be;  /* CLOCK_REALTIME seconds (network order) */
    uint32_t send_nsec_be; /* CLOCK_REALTIME nanoseconds (network order) */
} rb_hdr_t;
#pragma pack(pop)

/* --- utils --- */
static void error_exit(const char *msg) {
    perror(msg);
    exit(EXIT_FAILURE);
}

static void enable_hw_timestamping(int sockfd, const char *iface) {
    struct hwtstamp_config config = {0};
    struct ifreq ifr = {0};
    strncpy(ifr.ifr_name, iface, IFNAMSIZ);

    /* Request NIC RX hardware timestamping */
    config.flags = 0;
    config.tx_type = HWTSTAMP_TX_OFF;
    config.rx_filter = HWTSTAMP_FILTER_ALL;   // your NIC supports "all"

    ifr.ifr_data = (void *)&config;
    if (ioctl(sockfd, SIOCSHWTSTAMP, &ifr) < 0) {
        perror("ioctl SIOCSHWTSTAMP");
        exit(EXIT_FAILURE);
    }

    /* Choose the proper flag for ts[1] depending on your headers */
    #if defined(SOF_TIMESTAMPING_CLOCK_REALTIME)
        /* Newer kernels: ask for system-domain conversion in CLOCK_REALTIME */
        const int domain_flag = SOF_TIMESTAMPING_CLOCK_REALTIME;
    #elif defined(SOF_TIMESTAMPING_SYS_HARDWARE)
        /* Older kernels: this enables ts[1] in system domain */
        const int domain_flag = SOF_TIMESTAMPING_SYS_HARDWARE;
    #else
        const int domain_flag = 0;  /* (very old headers) */
    #endif

    /* Ask the socket to deliver ts[0], ts[1], ts[2] */
    int ts_flags =
        SOF_TIMESTAMPING_RX_HARDWARE |
        SOF_TIMESTAMPING_RAW_HARDWARE |
        SOF_TIMESTAMPING_RX_SOFTWARE |   /* <-- needed for ts[0] */
        SOF_TIMESTAMPING_SOFTWARE     |
        SOF_TIMESTAMPING_SYS_HARDWARE;                     /* <-- ts[1] domain select */

    if (setsockopt(sockfd, SOL_SOCKET, SO_TIMESTAMPING, &ts_flags, sizeof(ts_flags)) < 0) {
        perror("setsockopt SO_TIMESTAMPING");
        exit(EXIT_FAILURE);
    }
}

static void get_datetime_strings(char *date_str, size_t dsz, char *time_str, size_t tsz) {
    time_t now = time(NULL);
    struct tm *tm_info = localtime(&now);
    strftime(date_str, dsz, "%Y-%m-%d", tm_info);
    strftime(time_str, tsz, "%H-%M-%S", tm_info);
}

static char *create_log_folder(void) {
    static char base_path[256];
    char date_str[16], time_str[16];
    get_datetime_strings(date_str, sizeof(date_str), time_str, sizeof(time_str));
    snprintf(base_path, sizeof(base_path), "/tmp/tslog_%s_%s", date_str, time_str);
    if (mkdir(base_path, 0755) < 0 && errno != EEXIST) {
        error_exit("mkdir");
    }
    return base_path;
}

static FILE *create_log_file(const char *folder, const char *iface) {
    char file_path[512];
    snprintf(file_path, sizeof(file_path), "%s/%s.csv", folder, iface);
    FILE *f = fopen(file_path, "w");
    if (!f) error_exit("fopen");
    /* CSV header */
    fprintf(f, "seq,send_sec,send_nsec,sw_sec,sw_nsec,hw_sys_sec,hw_sys_nsec,hw_raw_sec,hw_raw_nsec,sys_sec,sys_nsec,recv_latency_ns\n");
    fflush(f);
    return f;
}

int main(int argc, char *argv[]) {
    /* realtime + mlock */
    struct sched_param param = { .sched_priority = 90 };
    if (sched_setscheduler(0, SCHED_FIFO, &param) < 0) error_exit("sched_setscheduler");
    if (mlockall(MCL_CURRENT | MCL_FUTURE) < 0) error_exit("mlockall");

    if (argc != 2) {
        fprintf(stderr, "Usage: %s <interface>\n", argv[0]);
        return 1;
    }
    const char *iface = argv[1];

    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) error_exit("socket");

    /* Bind to chosen interface */
    if (setsockopt(sockfd, SOL_SOCKET, SO_BINDTODEVICE, iface, strlen(iface)) < 0)
        error_exit("setsockopt SO_BINDTODEVICE");

    /* Big receive buffer to avoid drops */
    int rcvbuf = 8 * 1024 * 1024;
    setsockopt(sockfd, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof(rcvbuf)); /* best-effort */

    enable_hw_timestamping(sockfd, iface);

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port   = htons(RECV_PORT),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };
    if (bind(sockfd, (struct sockaddr *)&addr, sizeof(addr)) < 0)
        error_exit("bind");

    char *log_folder = create_log_folder();
    FILE *logfile = create_log_file(log_folder, iface);
    printf("Listening on %s and logging to %s/%s.csv\n", iface, log_folder, iface);

    /* I/O buffers */
    unsigned char data[MAXLINE];
    /* control buffer sized for 3 timespecs */
    char control[CMSG_SPACE(sizeof(struct timespec) * 3)];
    struct iovec iov = { .iov_base = data, .iov_len = sizeof(data) };
    struct msghdr msg;
    memset(&msg, 0, sizeof(msg));
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    while (1) {
        /* Reset control buffer length before each recvmsg */
        msg.msg_control = control;
        msg.msg_controllen = sizeof(control);
        msg.msg_name = NULL;
        msg.msg_namelen = 0;
        iov.iov_len = sizeof(data);

        struct timespec mono_before, mono_after;
        if (clock_gettime(CLOCK_MONOTONIC_RAW, &mono_before) < 0)
            perror("clock_gettime MONOTONIC_RAW before");

        ssize_t len = recvmsg(sockfd, &msg, 0);
        if (len < 0) {
            perror("recvmsg");
            continue;
        }

        if (clock_gettime(CLOCK_MONOTONIC_RAW, &mono_after) < 0)
            perror("clock_gettime MONOTONIC_RAW after");

        struct timespec sys_rt;
        if (clock_gettime(CLOCK_REALTIME, &sys_rt) < 0)
            perror("clock_gettime CLOCK_REALTIME");

        long recv_latency_ns = (mono_after.tv_sec - mono_before.tv_sec) * 1000000000L +
                               (mono_after.tv_nsec - mono_before.tv_nsec);

        /* Extract timestamps */
        // struct timespec ts_sw = {0}, ts_hw_sys = {0}, ts_hw_raw = {0};
        // struct cmsghdr *cmsg;
        // for (cmsg = CMSG_FIRSTHDR(&msg); cmsg != NULL; cmsg = CMSG_NXTHDR(&msg, cmsg)) {
        //     if (cmsg->cmsg_level == SOL_SOCKET && cmsg->cmsg_type == SO_TIMESTAMPING) {
        //         struct timespec *ts = (struct timespec *)CMSG_DATA(cmsg);
        //         /* ts[0]=software (CLOCK_REALTIME), ts[1]=transformed HW (CLOCK_REALTIME), ts[2]=raw HW (PHC) */
        //         ts_sw     = ts[0];
        //         ts_hw_sys = ts[1];
        //         ts_hw_raw = ts[2];
        //         break;
        //     }
        // }

        // inside the recv loop, before writing CSV:
        struct timespec ts_sw = {0}, ts_hw_sys = {0}, ts_hw_raw = {0};

        for (struct cmsghdr *c = CMSG_FIRSTHDR(&msg); c; c = CMSG_NXTHDR(&msg, c)) {
            if (c->cmsg_level == SOL_SOCKET &&
                c->cmsg_type  == SO_TIMESTAMPING &&
                c->cmsg_len   >= CMSG_LEN(sizeof(struct timespec) * 3)) {

                struct timespec *ts = (struct timespec *)CMSG_DATA(c);
                ts_sw     = ts[0];  // software (CLOCK_REALTIME)
                ts_hw_sys = ts[1];  // hardware -> system (CLOCK_REALTIME)
                ts_hw_raw = ts[2];  // raw PHC
                break;
            }
        }

        /* Parse RBRC header from payload */
        uint64_t seq = (uint64_t)-1;
        uint64_t send_sec = 0;
        uint32_t send_nsec = 0;
        if ((size_t)len >= sizeof(rb_hdr_t)) {
            rb_hdr_t hdr;
            memcpy(&hdr, data, sizeof(hdr));
            if (ntohl(hdr.magic) == MAGIC_RBRC && ntohl(hdr.version) == VERSION) {
                seq       = ntohll(hdr.seq_be);
                send_sec  = ntohll(hdr.send_sec_be);
                send_nsec = ntohl(hdr.send_nsec_be);
            } else {
                /* header did not match; keep seq = -1 to flag it */
            }
        }

        /* Write CSV row */
        fprintf(logfile,
                "%" PRIu64 ",%" PRIu64 ",%u,"
                "%ld,%ld,%ld,%ld,%ld,%ld,"
                "%ld,%ld,%ld\n",
                seq, send_sec, send_nsec,
                ts_sw.tv_sec,     ts_sw.tv_nsec,
                ts_hw_sys.tv_sec, ts_hw_sys.tv_nsec,
                ts_hw_raw.tv_sec, ts_hw_raw.tv_nsec,
                sys_rt.tv_sec, sys_rt.tv_nsec,
                recv_latency_ns);
        fflush(logfile);
    }

    fclose(logfile);
    close(sockfd);
    return 0;
}