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

#define RECV_PORT 5000
#define MAXLINE 1500

void error_exit(const char *msg) {
    perror(msg);
    exit(EXIT_FAILURE);
}

void enable_hw_timestamping(int sockfd, const char *iface) {
    struct hwtstamp_config config;
    struct ifreq ifr;
    memset(&ifr, 0, sizeof(ifr));
    strncpy(ifr.ifr_name, iface, IFNAMSIZ);

    config.flags = 0;
    config.tx_type = HWTSTAMP_TX_OFF;
    config.rx_filter = HWTSTAMP_FILTER_ALL;

    ifr.ifr_data = (void *)&config;

    if (ioctl(sockfd, SIOCSHWTSTAMP, &ifr) < 0) {
        error_exit("ioctl SIOCSHWTSTAMP");
    }

    int val = SOF_TIMESTAMPING_RX_HARDWARE |
              SOF_TIMESTAMPING_RAW_HARDWARE |
              SOF_TIMESTAMPING_SOFTWARE;

    if (setsockopt(sockfd, SOL_SOCKET, SO_TIMESTAMPING, &val, sizeof(val)) < 0)
        error_exit("setsockopt SO_TIMESTAMPING");
}

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);
}

char *create_log_folder(void) {
    static char base_path[128];
    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) {
        perror("mkdir");
        exit(EXIT_FAILURE);
    }

    return base_path;
}

FILE *create_log_file(const char *folder, const char *iface) {
    char file_path[256];
    snprintf(file_path, sizeof(file_path), "%s/%s.csv", folder, iface);

    FILE *f = fopen(file_path, "w");
    if (!f) {
        perror("fopen");
        exit(EXIT_FAILURE);
    }

    fprintf(f, "seq,rx_sec,rx_nsec,sys_sec,sys_nsec,elapsed_nsec\n");
    fflush(f);
    return f;
}

int main(int argc, char *argv[]) {
    struct sched_param param;
    param.sched_priority = 90;

    if (sched_setscheduler(0, SCHED_FIFO, &param) < 0) {
        perror("sched_setscheduler");
        exit(EXIT_FAILURE);
    } else {
        printf("Running with real-time priority (SCHED_FIFO, priority 90)\n");
    }

    if (mlockall(MCL_CURRENT | MCL_FUTURE) < 0) {
        perror("mlockall");
        exit(EXIT_FAILURE);
    }

    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");

    if (setsockopt(sockfd, SOL_SOCKET, SO_BINDTODEVICE, iface, strlen(iface)) < 0)
        error_exit("setsockopt SO_BINDTODEVICE");

    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);

    char data[MAXLINE];
    char control[512];
    struct iovec iov = { .iov_base = data, .iov_len = sizeof(data) };
    struct msghdr msg = {
        .msg_name = NULL,
        .msg_namelen = 0,
        .msg_iov = &iov,
        .msg_iovlen = 1,
        .msg_control = control,
        .msg_controllen = sizeof(control)
    };

    unsigned long seq = 0;

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

        struct timespec *hwts = NULL;
        struct timespec sys_ts_before, sys_ts_after;

        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);
                hwts = &ts[2];  // hardware timestamp
            }
        }

        if (clock_gettime(CLOCK_MONOTONIC_RAW, &sys_ts_before) < 0)
            perror("clock_gettime MONOTONIC_RAW before");

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

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

        long elapsed_nsec = (sys_ts_after.tv_sec - sys_ts_before.tv_sec) * 1e9 +
                            (sys_ts_after.tv_nsec - sys_ts_before.tv_nsec);

        if (hwts) {
            fprintf(logfile, "%lu,%ld,%ld,%ld,%ld,%ld\n", seq++, hwts->tv_sec, hwts->tv_nsec,
                    sys_ts.tv_sec, sys_ts.tv_nsec, elapsed_nsec);
        } else {
            fprintf(logfile, "%lu,-1,-1,%ld,%ld,%ld\n", seq++, sys_ts.tv_sec, sys_ts.tv_nsec, elapsed_nsec);
        }

        fflush(logfile);
    }

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