// gcc fastfetch_server.c -o fastfetch_server -O3 -lpthread
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <time.h>
#include <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

typedef struct {
    uint8_t addr[16];
    int count;
    double window_start;
} ratelimit_t;

static struct {
    ratelimit_t data[256];
    size_t next;
    pthread_mutex_t lock;
} ratelimit_table = { {0}, 0, PTHREAD_MUTEX_INITIALIZER };

typedef struct pthread_input_t {
    int client_fd;
    float window;
    int rate_limit;
    const char* command;
} pthread_input_t;

static double now(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec + ts.tv_nsec * 1e-9;
}

static int rate_limited(const uint8_t addr[16], float window, int rate_limit) {
    double n = now();

    pthread_mutex_lock(&ratelimit_table.lock);

    ratelimit_t* e = NULL;
    for (size_t i = 0; i < 256; i++) {
        if (memcmp(ratelimit_table.data[i].addr, addr, 16) == 0) {
            e = &ratelimit_table.data[i];
            break;
        }
    }

    if (!e) {
        e = &ratelimit_table.data[ratelimit_table.next++];
        memcpy(e->addr, addr, 16);
        e->window_start = n;
        e->count = 0;
    }

    if (n - e->window_start >= window) {
        e->window_start = n;
        e->count = 0;
    }

    const int allowed = e->count < rate_limit;
    if (allowed) e->count++;

    pthread_mutex_unlock(&ratelimit_table.lock);
    return allowed ? 0 : 1;
}

static void* handle_client(void* _arg) {
    pthread_input_t* arg = _arg;
    int client_fd = arg->client_fd;
    float window = arg->window;
    int rate_limit = arg->rate_limit;
    const char* command = arg->command;

    struct sockaddr_in6 peer;
    socklen_t peerlen = sizeof(peer);
    getpeername(client_fd, (struct sockaddr*)&peer, &peerlen);

    if (rate_limited(peer.sin6_addr.s6_addr, window, rate_limit)) {
        dprintf(client_fd,
            "HTTP/1.1 429 Too Many Requests\r\n"
            "Content-Type: text/plain\r\n"
            "Content-Length: 19\r\n"
            "Connection: close\r\n"
            "\r\n"
            "Too Many Requests\r\n"
        );
        goto cleanup;
    }

    char fastfetch[16384];
    FILE* fp = popen(command, "r");
    size_t n = fp ? fread(fastfetch, sizeof(fastfetch[0]), sizeof(fastfetch) - 1, fp) : 0;
    if (fp) pclose(fp);
    fastfetch[n] = '\0';

    dprintf(client_fd,
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/plain\r\n"
        "Content-Length: %zu\r\n"
        "Connection: close\r\n"
        "\r\n"
        "%s",
        n, fastfetch);

cleanup:
    shutdown(client_fd, SHUT_WR);
    char discard[1024];
    while (recv(client_fd, discard, sizeof(discard), 0) > 0) {}
    close(client_fd);
    return NULL;
}

#define streq(a, b) (strcmp((a), (b)) == 0)
#define exit_error(...) do { \
    fprintf(stderr, __VA_ARGS__); \
    exit(EXIT_FAILURE); \
} while (0)

#define exit_perror(s) do { \
    perror(s); \
    exit(EXIT_FAILURE); \
} while (0)

int main(int argc, char *argv[]) {
    int port = 80;
    float window = 60.0f;
    int rate_limit = 5;
    const char* command = "fastfetch --pipe false 2>/dev/null";
    struct in6_addr bind_addr = in6addr_any;

    for (int i = 1; i < argc; i++) {
        if (streq(argv[i], "--help")) {
            puts("usage:\n"
                 "\tfastfetch_server --addr 127.0.01 (optional)\n"
                 "\t                 --window 60 (optional)\n"
                 "\t                 --ratelimit 5 (optional)\n"
                 "\t                 --hyfetch (optional)\n"
                 "\t                 port (optional)\n"
                 "\t --addr expects an address in an IPv4 or IPv6 format\n"
                 "\t --window sets a rate limit delay in seconds\n"
                 "\t --ratelimit sets the rate limit amount\n"
                 "\t --hyfetch uses hyfetch with fastfetch backend instead of fastfetch\n"
                 "\t port specifies the port, 80 by default, might need superuser for\n"
                 "\t      ports below 1024\n"
                 "get latest version with:\n"
                 "\t wget https://gitea.codersquack.nl/tema5002/.profile/raw/branch/main/fastfetch_server.c\n"
                 "\t gcc fastfetch_server.c -o fastfetch_server -O3 -lpthread\n"
                 "so far tested on:\n"
                 "\t GNU/Linux on x86_32/x86_64 with glibc/musl");
            return 0;
        }
        if (streq(argv[i], "--addr") && i + 1 < argc) {
            if (inet_pton(AF_INET6, argv[++i], &bind_addr) != 1) {
                struct in_addr v4;
                if (inet_pton(AF_INET, argv[i], &v4) != 1) exit_error("invalid address: %s\n", argv[i]);
                memset(&bind_addr, 0, sizeof(bind_addr));
                bind_addr.s6_addr[10] = 0xff;
                bind_addr.s6_addr[11] = 0xff;
                memcpy(&bind_addr.s6_addr[12], &v4, 4);
            }
        }
        else if (streq(argv[i], "--window") && i + 1 < argc) {
            window = atof(argv[++i]);
        }
        else if (streq(argv[i], "--ratelimit") && i + 1 < argc) {
            rate_limit = atoi(argv[++i]);
        }
        else if (streq(argv[i], "--hyfetch")) {
            command = "hyfetch --backend=fastfetch --args=\"--pipe false\" 2>/dev/null";
        }
        else {
            port = atoi(argv[i]);
            if (port <= 0) exit_error("unknown argument: %s\n", argv[i]);
        }
    }

    const int server_fd = socket(AF_INET6, SOCK_STREAM, 0);
    if (server_fd < 0) exit_perror("socket");

    int opt = 0;
    setsockopt(server_fd, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt));
    opt = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    const struct sockaddr_in6 addr = {
        .sin6_family = AF_INET6,
        .sin6_addr   = bind_addr,
        .sin6_port   = htons(port)
    };

    if (bind(server_fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) exit_perror("bind");
    if (listen(server_fd, 128) < 0) exit_perror("listen");

    char addr_str[INET6_ADDRSTRLEN];
    inet_ntop(AF_INET6, &bind_addr, addr_str, sizeof(addr_str));
    printf("listening on [%s]:%d  (limit: %d req / %fs per IP)\n",
           addr_str, port, rate_limit, window);

    for (;;) {
        struct sockaddr_in6 client_addr;
        socklen_t addrlen = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &addrlen);
        if (client_fd < 0) continue;

        pthread_input_t pi = {.client_fd = client_fd, .window = window, .rate_limit = rate_limit, .command = command};

        pthread_t tid;
        if (pthread_create(&tid, NULL, handle_client, (void*)(intptr_t)&pi) == 0)
            pthread_detach(tid);
        else {
            perror("pthread_create");
            close(client_fd);
        }
    }
}