#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define VTERM_ROWS 50
#define VTERM_COLS 200
#include "ansi2html.h"

#include "stdint.h"
#define RATELIMIT_THREADSAFE
typedef uint8_t ratelimit_table_type[16];
#include "ratelimit.h"

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

#define CURL_HEADERS(status) \
    "HTTP/1.1 " status "\r\n" \
    "Content-Type: text/plain\r\n" \
    "Content-Length: %zu\r\n" \
    "Connection: close\r\n" \
    "\r\n" \
    "%s"

#define BROWSER_HEADERS(status) \
    "HTTP/1.1 " status "\r\n" \
    "Content-Type: text/html; charset=utf-8\r\n" \
    "Content-Length: %zu\r\n" \
    "Connection: close\r\n" \
    "\r\n" \
    "%s"

#define FS_HELP \
    "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" \
     "so far tested on:\n" \
     "\t GNU/Linux on x86/x86_64 with glibc/musl\n"

static void send_fastfetch(int client_fd, const char* command, int is_browser) {
    char fastfetch[16384];
    FILE* fp = popen(command, "r");
    const size_t n = fp ? fread(fastfetch, sizeof(fastfetch[0]), sizeof(fastfetch) - 1, fp) : 0;
    if (fp) pclose(fp);
    fastfetch[n] = '\0';

    if (!is_browser) {
        dprintf(client_fd, CURL_HEADERS("200 OK"), n, fastfetch);
    }
    else {
        char* html = ansi2html(fastfetch);
        dprintf(client_fd, BROWSER_HEADERS("200 OK"), strlen(html), html);
        free(html);
    }
}

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, CURL_HEADERS("429 Too Many Requests"), 19, "Too Many Requests\r\n");
        goto cleanup;
    }

    char what_to_do[256] = {0};
    int is_browser = 0;
    {
        struct {
            char buf[8192];
            size_t len;
        } req = {.len = 0};

        while (req.len < sizeof(req.buf) - 1) {
            const ssize_t n = recv(client_fd, req.buf + req.len, sizeof(req.buf) - 1 - req.len, 0);
            if (n <= 0) break;
            req.len += n;
            req.buf[req.len] = '\0';
            const char* body = memmem(req.buf, req.len, "\r\n\r\n", 4);

            if (body) {
                for (const char* p = req.buf; p < body;) {
                    const char* nl = memchr(p, '\n', body - p);
                    const size_t line_len = nl ? (size_t)(nl - p) : (size_t)(body - p);

                    if (line_len >= 14 && strncasecmp(p, "get /", 5) == 0) {
                        const char* protocol = memmem(p, line_len, " HTTP/1.1", 9);
                        if (!protocol) goto continue_label;
                        const size_t get_str_len = protocol - p - 5;
                        if (get_str_len >= sizeof(what_to_do)) goto cleanup;
                        memcpy(what_to_do, p + 5, get_str_len);
                    }
                    else if (line_len >= 12 && strncasecmp(p, "user-agent:", 11) == 0) {
                        const char* ua = p + 11;
                        while (*ua == ' ' || *ua == '\t') ua++;
                        is_browser = strncmp(ua, "Mozilla", 7) == 0 && (
                            memmem(ua, line_len, "Chrome", 6) ||
                            memmem(ua, line_len, "Firefox", 7) ||
                            memmem(ua, line_len, "Safari", 6));
                        goto we_now_know_if_we_got_a_browser;
                    }
continue_label:
                    p = nl ? nl + 1 : body;
                }
            }
        }
    }
we_now_know_if_we_got_a_browser:
    if (what_to_do[0] == '\0') {
        send_fastfetch(client_fd, command, is_browser);
        goto cleanup;
    }
    if (strcmp(what_to_do, "help") == 0) {
        dprintf(client_fd, CURL_HEADERS("200 OK"), sizeof(FS_HELP) - 1, FS_HELP);
        goto cleanup;
    }
    dprintf(client_fd, CURL_HEADERS("404 Not Found"), 11, "Not Found\r\n");

cleanup:
    shutdown(client_fd, SHUT_WR);
    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")) {
            printf(FS_HELP);
            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);
        }
    }
}