#pragma once

#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <pthread.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/socket.h>
#include <unistd.h>
//
#include "common.c"
#include "ratelimit.c"

typedef uint8_t byte;
typedef void (*tome_ptr)(
    char *ip_addr, size_t addr_size, byte *in, size_t in_len, byte **out, size_t *out_len);

#define MAX_BUFFER      1 << 18
#define MAX_BODY        1 << 22
#define MAX_CONNECTIONS 2048

// vars

int worker_count = 64;

pthread_mutex_t queue_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t  queue_cond  = PTHREAD_COND_INITIALIZER;

static int    ring_fds[ MAX_CONNECTIONS ];
static char   ring_ips[ MAX_CONNECTIONS ][ INET6_ADDRSTRLEN ];
static size_t ring_head  = 0;
static size_t ring_tail  = 0;
size_t        queue_size = 0;

tome_ptr global_handler = NULL;

// utils

static ssize_t write_all(int fd, const void *buf, size_t len) {
    size_t      off = 0;
    const char *p   = buf;
    while (off < len) {
        ssize_t r = send(fd, p + off, len - off, MSG_NOSIGNAL);
        if (r < 0) {
            if (errno == EINTR) continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK) continue;
            return -1;
        }
        off += (size_t) r;
    }
    return (ssize_t) off;
}

static int read_until_headers(int fd, char **out_buf, size_t *out_len) {
    size_t total_tries = 0;

    size_t cap = 4096;
    size_t len = 0;
    char  *buf = malloc(cap);
    if (!buf) return -1;
    while (1) {
        if (total_tries++ > 128) break;

        if (cap - len < 1024) {
            size_t ncap = cap * 2;
            char  *n    = realloc(buf, ncap);
            if (!n) {
                free(buf);
                return -1;
            }
            buf = n;
            cap = ncap;
        }
        ssize_t r = recv(fd, buf + len, cap - len, 0);
        if (r < 0) {
            if (errno == EINTR) continue;
            free(buf);
            return -1;
        } else if (r == 0) {
            // TODO: why?
            // len = len;
            break;
        }
        len += (size_t) r;
        if (len >= 4 && memmem(buf, len, "\r\n\r\n", 4) != NULL) {
            if (len >= 19 && memmem(buf, len, "Transfer-Encoding: chunked", 19) != NULL) {
                const char *err = "HTTP/1.1 501 Not Implemented\r\n"
                                  "Content-Length: 31\r\n"
                                  "Connection: close\r\n\r\n"
                                  "chunked encoding not supported";
                write_all(fd, err, strlen(err));
                free(buf);
                return -1;
            }
            break;
        }
        if (len > MAX_BUFFER) break;
    }
    *out_buf = buf;
    *out_len = len;
    return 0;
}

static int read_exact(int fd, void *buf, size_t count) {
    size_t off = 0;
    char  *p   = buf;
    while (off < count) {
        ssize_t r = recv(fd, p + off, count - off, 0);
        if (r < 0) {
            if (errno == EINTR) continue;
            return -1;
        } else if (r == 0) {
            return -1;
        }
        off += (size_t) r;
    }
    return 0;
}

static ssize_t parse_content_length(const char *hdr, size_t hdr_len) {
    const char *p   = hdr;
    const char *end = hdr + hdr_len;
    while (p < end) {
        const char *line_end = memchr(p, '\n', (size_t) (end - p));
        if (!line_end) line_end = end;
        size_t line_len = (size_t) (line_end - p);
        if (line_len >= 15) {
            if (strncasecmp(p, "Content-Length:", 15) == 0) {
                const char *num = p + 15;
                while (num < line_end && (*num == ' ' || *num == '\t')) num++;
                char   tmp[ 32 ] = { 0 };
                size_t to_copy   = line_len - (num - p);
                if (to_copy >= sizeof(tmp)) to_copy = sizeof(tmp) - 1;
                memcpy(tmp, num, to_copy);
                for (int i = (int) strlen(tmp) - 1; i >= 0
                                                    && (tmp[ i ] == '\r' || tmp[ i ] == '\n'
                                                        || tmp[ i ] == ' ' || tmp[ i ] == '\t');
                     --i)
                    tmp[ i ] = 0;
                long v = strtol(tmp, NULL, 10);
                if (v < 0) return -1;
                return (ssize_t) v;
            }
        }
        p = line_end + 1;
    }
    return -1;
}

static int send_http_response(int fd, const byte *payload, size_t payload_len) {
    const char *template_hdr
        = "HTTP/1.1 200 OK\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n";
    char hdr[ 128 ];
    int  n = snprintf(hdr, sizeof(hdr), template_hdr, payload_len);
    if (n < 0 || (size_t) n >= sizeof(hdr)) return -1;
    if (write_all(fd, hdr, (size_t) n) < 0) return -1;
    if (payload_len > 0 && payload) {
        if (write_all(fd, payload, payload_len) < 0) return -1;
    }
    return 0;
}

// server

void *omen_bound(void *arg) {
    (void) arg;
    while (1) {
        int  client_fd                     = -1;
        char client_ip[ INET6_ADDRSTRLEN ] = { 0 };

        pthread_mutex_lock(&queue_mutex);
        while (queue_size == 0) { pthread_cond_wait(&queue_cond, &queue_mutex); }

        client_fd = ring_fds[ ring_head ];
        strncpy(client_ip, ring_ips[ ring_head ], INET6_ADDRSTRLEN - 1);
        client_ip[ INET6_ADDRSTRLEN - 1 ] = '\0';
        ring_head                         = (ring_head + 1) % MAX_CONNECTIONS;
        queue_size--;
        pthread_mutex_unlock(&queue_mutex);

        struct timeval timeout = { 5, 0 };
        setsockopt(client_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

        char  *hdr_buf = NULL;
        size_t hdr_len = 0;
        if (read_until_headers(client_fd, &hdr_buf, &hdr_len) < 0) {
            close(client_fd);
            free(hdr_buf);
            continue;
        }

        char  *sep      = memmem(hdr_buf, hdr_len, "\r\n\r\n", 4);
        size_t body_len = 0;
        byte  *body_ptr = NULL;

        if (sep) {
            size_t  headers_bytes = (size_t) (sep - hdr_buf) + 4;
            ssize_t cl            = parse_content_length(hdr_buf, headers_bytes);
            if (cl > 0) {
                size_t already = hdr_len - headers_bytes;
                if ((size_t) cl > MAX_BODY) cl = MAX_BODY;
                body_ptr = malloc((size_t) cl);
                if (!body_ptr) {
                    free(hdr_buf);
                    close(client_fd);
                    continue;
                }
                if (already > 0) {
                    size_t copy = already;
                    if (copy > (size_t) cl) copy = (size_t) cl;
                    memcpy(body_ptr, sep + 4, copy);
                    body_len = copy;
                }
                if ((size_t) body_len < (size_t) cl) {
                    if (read_exact(client_fd, body_ptr + body_len, (size_t) cl - body_len) < 0) {
                        free(body_ptr);
                        free(hdr_buf);
                        close(client_fd);
                        continue;
                    }
                    body_len = (size_t) cl;
                }
            } else {
                size_t already = hdr_len - (size_t) (sep - hdr_buf) - 4;
                if (already > 0) {
                    if (already > MAX_BODY) already = MAX_BODY;
                    body_ptr = malloc(already);
                    if (body_ptr) {
                        memcpy(body_ptr, sep + 4, already);
                        body_len = already;
                    }
                } else {
                    body_ptr = NULL;
                    body_len = 0;
                }
            }
        } else {
            size_t use = hdr_len;
            if (use > MAX_BODY) use = MAX_BODY;
            body_ptr = malloc(use);
            if (body_ptr) {
                memcpy(body_ptr, hdr_buf, use);
                body_len = use;
            }
        }

        byte  *out     = NULL;
        size_t out_len = 0;

        if (global_handler)
            global_handler(client_ip, sizeof(client_ip), body_ptr, body_len, &out, &out_len);
        if (out_len > MAX_BODY) out_len = MAX_BODY;

        send_http_response(client_fd, out, out_len);

        free(out);
        free(body_ptr);
        free(hdr_buf);
        close(client_fd);
    }
    return NULL;
}

int omen(int port, tome_ptr handler) {
    global_handler = handler;
    signal(SIGPIPE, SIG_IGN);

    int server_fd = socket(AF_INET6, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("failed to create IPv6 socket");
        return -1;
    }

    int opt = 0;
    if (setsockopt(server_fd, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt)) < 0) {
        perror("failed to disable IPV6_V6ONLY");
        close(server_fd);
        return -1;
    }

    opt = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        perror("failed to set SO_REUSEADDR");
        close(server_fd);
        return -1;
    }

    struct sockaddr_in6 addr = { 0 };
    addr.sin6_family         = AF_INET6;
    addr.sin6_addr           = in6addr_any;
    addr.sin6_port           = htons(port);

    if (bind(server_fd, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
        perror("bind failed");
        close(server_fd);
        return -1;
    }

    if (listen(server_fd, MAX_CONNECTIONS) < 0) {
        perror("listen failed");
        close(server_fd);
        return -1;
    }

    pthread_t *threads = malloc((size_t) worker_count * sizeof(pthread_t));
    if (!threads) {
        close(server_fd);
        return -1;
    }

    for (int i = 0; i < worker_count; i++) {
        if (pthread_create(&threads[ i ], NULL, omen_bound, NULL) != 0) {
            perror("pthread_create failed");
            free(threads);
            close(server_fd);
            return -1;
        }
        pthread_detach(threads[ i ]);
    }

    while (1) {
        struct sockaddr_storage client_addr;
        socklen_t               len = sizeof(client_addr);
        int client_fd               = accept(server_fd, (struct sockaddr *) &client_addr, &len);
        if (client_fd < 0) {
            if (errno != EINTR) { perror("accept failed"); }
            continue;
        }

        char client_ip[ INET6_ADDRSTRLEN ] = { 0 };
        if (client_addr.ss_family == AF_INET) {
            struct sockaddr_in *addr4 = (struct sockaddr_in *) &client_addr;
            inet_ntop(AF_INET, &addr4->sin_addr, client_ip, sizeof(client_ip));
        } else if (client_addr.ss_family == AF_INET6) {
            struct sockaddr_in6 *addr6 = (struct sockaddr_in6 *) &client_addr;
            if (IN6_IS_ADDR_V4MAPPED(&addr6->sin6_addr)) {
                struct in_addr v4;
                memcpy(&v4.s_addr, &addr6->sin6_addr.s6_addr[ 12 ], sizeof(v4.s_addr));
                inet_ntop(AF_INET, &v4, client_ip, sizeof(client_ip));
            } else {
                inet_ntop(AF_INET6, &addr6->sin6_addr, client_ip, sizeof(client_ip));
            }
        } else {
            strncpy(client_ip, "0.0.0.0", sizeof(client_ip) - 1);
            client_ip[ sizeof(client_ip) - 1 ] = '\0';
        }

        if (!ratelimit_check(client_ip)) { usleep(250 * 1000); }

        pthread_mutex_lock(&queue_mutex);
        if (queue_size >= (size_t) MAX_CONNECTIONS) {
            char err_resp[] = "HTTP/1.1 418 I'm a teapot\r\n"
                              "Content-Length: 24\r\n"
                              "Connection: close\r\n\r\n"
                              "Max connections reached";
            write_all(client_fd, err_resp, strlen(err_resp));
            close(client_fd);
            pthread_mutex_unlock(&queue_mutex);
            continue;
        }

        ring_fds[ ring_tail ] = client_fd;
        strncpy(ring_ips[ ring_tail ], client_ip, INET6_ADDRSTRLEN - 1);
        ring_ips[ ring_tail ][ INET6_ADDRSTRLEN - 1 ] = '\0';
        ring_tail                                     = (ring_tail + 1) % MAX_CONNECTIONS;
        queue_size++;
        pthread_cond_signal(&queue_cond);
        pthread_mutex_unlock(&queue_mutex);
    }

    free(threads);
    close(server_fd);
    return 0;
}

// client

int study(
    const char *ip,
    int         port,
    const byte *data,
    size_t      data_len,
    byte      **response,
    size_t     *resp_len) {
    bool is_ipv6 = strchr(ip, ':') != NULL;
    int  domain  = is_ipv6 ? AF_INET6 : AF_INET;
    // printf("attempting %s connection to %s:%d\n", is_ipv6 ? "ipv6" : "ipv4", ip, port);

    int sock = socket(domain, SOCK_STREAM, 0);
    if (sock < 0) {
        printf("socket creation failed: %s\n", strerror(errno));
        return -1;
    }

    struct sockaddr_storage addr = { 0 };
    addr.ss_family               = domain;
    void  *addr_ptr              = NULL;
    size_t addr_len              = 0;
    if (domain == AF_INET) {
        struct sockaddr_in *addr4 = (struct sockaddr_in *) &addr;
        addr_ptr                  = &addr4->sin_addr;
        addr_len                  = sizeof(struct sockaddr_in);
        addr4->sin_port           = htons(port);
    } else {
        struct sockaddr_in6 *addr6 = (struct sockaddr_in6 *) &addr;
        addr_ptr                   = &addr6->sin6_addr;
        addr_len                   = sizeof(struct sockaddr_in6);
        addr6->sin6_port           = htons(port);
    }

    if (inet_pton(domain, ip, addr_ptr) != 1) {
        printf("invalid %s address: %s\n", is_ipv6 ? "ipv6" : "ipv4", ip);
        close(sock);
        return -1;
    }

    struct timeval timeout = { 5, 0 };
    if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0
        || setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0) {
        printf("setsockopt failed: %s\n", strerror(errno));
        close(sock);
        return -1;
    }

    // printf("connecting to %s:%d\n", ip, port);
    if (connect(sock, (struct sockaddr *) &addr, addr_len) < 0) {
        printf("connect failed: %s\n", strerror(errno));
        close(sock);
        return -1;
    }

    if (data_len > MAX_BODY) data_len = MAX_BODY;
    char req_hdr[ 256 ];
    int  hdr_n;
    if (is_ipv6) {
        hdr_n = snprintf(
            req_hdr,
            sizeof(req_hdr),
            "POST /data HTTP/1.1\r\nHost: [%s]\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n",
            ip,
            data_len);
    } else {
        hdr_n = snprintf(
            req_hdr,
            sizeof(req_hdr),
            "POST /data HTTP/1.1\r\nHost: %s\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n",
            ip,
            data_len);
    }
    if (hdr_n < 0 || (size_t) hdr_n >= sizeof(req_hdr)) {
        printf("failed to format HTTP header\n");
        close(sock);
        return -1;
    }

    // printf("sending request: %.*s\n", hdr_n, req_hdr);
    if (write_all(sock, req_hdr, (size_t) hdr_n) < 0
        || (data_len > 0 && write_all(sock, data, data_len) < 0)) {
        printf("write failed: %s\n", strerror(errno));
        close(sock);
        return -1;
    }

    char  *hdr_buf = NULL;
    size_t hdr_len = 0;
    if (read_until_headers(sock, &hdr_buf, &hdr_len) < 0) {
        printf("failed to read headers: %s\n", strerror(errno));
        close(sock);
        free(hdr_buf);
        return -1;
    }

    if (memmem(hdr_buf, hdr_len, "HTTP/1.1 418 I'm a teapot", 34) != NULL) {
        // printf("server responded with 418 I'm a teapot\n");
        close(sock);
        free(hdr_buf);
        return -2;
    }

    char  *sep      = memmem(hdr_buf, hdr_len, "\r\n\r\n", 4);
    size_t body_len = 0;
    byte  *body     = NULL;
    if (sep) {
        size_t  headers_bytes = (size_t) (sep - hdr_buf) + 4;
        ssize_t cl            = parse_content_length(hdr_buf, headers_bytes);
        size_t  already       = hdr_len - headers_bytes;
        if (cl > 0) {
            if ((size_t) cl > MAX_BODY) cl = MAX_BODY;
            body = malloc((size_t) cl);
            if (!body) {
                printf("failed to allocate body buffer\n");
                free(hdr_buf);
                close(sock);
                return -1;
            }
            if (already > 0) {
                size_t copy = already;
                if (copy > (size_t) cl) copy = (size_t) cl;
                memcpy(body, sep + 4, copy);
                body_len = copy;
            }
            if ((size_t) body_len < (size_t) cl) {
                if (read_exact(sock, body + body_len, (size_t) cl - body_len) < 0) {
                    printf("failed to read body: %s\n", strerror(errno));
                    free(body);
                    free(hdr_buf);
                    close(sock);
                    return -1;
                }
                body_len = (size_t) cl;
            }
        } else {
            if (already > 0) {
                if (already > MAX_BODY) already = MAX_BODY;
                body = malloc(already);
                if (body) {
                    memcpy(body, sep + 4, already);
                    body_len = already;
                }
            } else {
                body     = NULL;
                body_len = 0;
            }
        }
    } else {
        size_t use = hdr_len;
        if (use > MAX_BODY) use = MAX_BODY;
        body = malloc(use);
        if (body) {
            memcpy(body, hdr_buf, use);
            body_len = use;
        }
    }

    free(hdr_buf);
    close(sock);

    if (body_len > 0 && body) {
        *response = body;
        *resp_len = body_len;
    } else {
        *response = NULL;
        *resp_len = 0;
        free(body);
    }
    // printf("response received, length: %zu\n", body_len);
    return 0;
}