#include <assert.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <pthread.h>
#include <wait.h>

#define ALLOWED_PATTERNS_FILE "allowed_patterns.txt"
#define ALLOWED_IPS_FILE "allowed_ips.txt"
#define PATTERNS_FILE "patterns.txt"
// #define ACCESS_LOG_FILE "access.log.1"
#define ACCESS_LOG_FILE "/var/log/nginx/access.log"

#define print_bool_opt(x) \
    if (x) { \
        fprintf(stderr, #x ": true\n"); \
    } else { \
        fprintf(stderr, #x ": false\n"); \
    }
#define print_str_opt(x) \
    if (x != nullptr) { \
        fprintf(stderr, #x ": %s\n", x); \
    } else { \
        fprintf(stderr, #x ": null\n"); \
    }

#define ansi(x) "\033[" x
#define ansi_mod(x) "\033[" x "m"

#define ANSI_CLEAR_LINE ansi("2K")
#define ANSI_BLINK ansi_mod("5")
#define ANSI_RESET ansi_mod("0")

#define ANSI_COLOR(r,g,b) ansi_mod("38;2;" #r ";" #g ";" #b)
#define ANSI_HSL(h,s,l) ansi_mod("38;2;" \
    (int)((l < 50 ? l + s : l - s) * (1 + h / 360.0) / 100.0 * 255) ";" \
    (int)((l < 50 ? l + s : l - s) * (1 - h / 360.0) / 100.0 * 255) ";" \
    (int)(l * 255 / 100))

typedef struct
{
    char** patterns;
    char** allowed_patterns;
    char** allowed_ips;
    // size_t allowedIpCount;
    char** banned_ips;
    int* banned_ip_hits;
    size_t banned_ip_count;
    size_t banned_ip_size;
} MatchRules;

bool print_allowed_ip = false;
bool print_allowed_pattern = false;
bool print_already_banned_ip = false;
bool print_indeterminate = false;
bool print_bans = true;
bool print_only = false;
bool print_progress = true;

bool check_allowed_ip = true;
bool check_allowed_pattern = true;
bool check_already_banned_ip = true;

bool dump_patterns = false;

char* allowed_patterns_path = ALLOWED_PATTERNS_FILE;
char* allowed_ips_path = ALLOWED_IPS_FILE;
char* patterns_path = PATTERNS_FILE;
char* access_log_path = ACCESS_LOG_FILE;

void* xmalloc(size_t size)
{
    void* ptr = malloc(size);
    if (ptr == NULL)
    {
        fprintf(stderr, "Memory allocation failed\n");
        exit(EXIT_FAILURE);
    }
    return ptr;
}

char* str_to_hex(const char* str)
{
    const size_t len = strlen(str);
    char* hex = xmalloc(len * 2 + 3);
    for (size_t i = 0; i < len; i++)
    {
        sprintf(hex + i * 3, "%02x ", (unsigned char)str[i]);
    }
    sprintf(hex + len * 3, "00 ");
    hex[len * 3 + 2] = '\0';
    return hex;
}

bool streq(const char* a, const char* b)
{
    return a == b || (a != nullptr && b != nullptr && strcmp(a, b) == 0);
}

void signal_handler(int signal_number)
{
    int wait_status;
    pid_t return_pid = wait(&wait_status);
    if (return_pid == -1)
    {
        perror("wait()");
    }
    if (WIFEXITED(wait_status))
    {
        printf("job [%d] | pid: %d | exit status: %d\n", signal_number, return_pid, WEXITSTATUS(wait_status));
    }
    else
    {
        printf("exit abnormally\n");
    }

    fprintf(stderr, "the signal %d was received\n", signal_number);
}

void read_lines_cb(const char* filename, void cb(char*, void*), void* cbData)
{
    if (cb == nullptr)
    {
        fprintf(stderr, "Callback function is null\n");
        return;
    }

    clock_t ctime = clock(), lastPrint = clock();
    FILE* file = fopen(filename, "r");
    if (file == NULL)
    {
        fprintf(stderr, "Error opening file %s: ", filename);
        perror("");
        return;
    }

    int count = 0;
    char* buf = xmalloc(10000);
    clock_t times[1000];
    while (fgets(buf, 10000, file) != NULL)
    {
        if (buf[0] == '#')
        {
            continue;
        }
        count++;
        buf[strcspn(buf, "\n")] = 0;
        char* line = strdup(buf);

        cb(line, cbData);
        free(line);

        times[count % 1000] = clock();
        if (print_progress && count % 10 == 0)
        {
            double elapsed = (clock() - lastPrint) * 1000.0 / CLOCKS_PER_SEC;
            double elapsed_total = (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC;
            if (elapsed > 50.0)
            {
                fprintf(stderr, "\rRead %d lines from %s in %.3f ms (%.3f/s)\r", count, filename, elapsed_total, count / (elapsed_total / 1000.0));
                lastPrint = clock();
            }
        }
    }

    free(buf);

    fclose(file);
    fprintf(stderr, "Read %d definitions from %s in %.3f ms (%.3f/s)\n", count, filename, (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC, count * 1000.0 / (clock() - ctime));
}

char** read_lines(const char* filename)
{
    clock_t ctime = clock(), lastPrint = clock();
    FILE* file = fopen(filename, "r");
    if (file == NULL)
    {
        fprintf(stderr, "Error opening file %s: ", filename);
        perror("");
        return nullptr;
    }

    int count = 0;
    {
        char* buf = xmalloc(10000);
        while (fgets(buf, 10000, file) != NULL)
        {
            if (buf[0] == '#')
            {
                continue;
            }
            count++;
        }

        free(buf);
        rewind(file);
    }
    fprintf(stderr, "Counted %d lines from %s in %.3f ms\n", count, filename, (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC);

    // count = count_lines(file);
    // printf("Counted %d lines from %s in %.3f ms\n", count, filename, (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC);


    rewind(file);
    char** lines = xmalloc((count + 1) * sizeof(char*));

    char* buf = xmalloc(10000);
    for (int i = 0; i < count; i++)
    {
        if (fgets(buf, 10000, file) == NULL)
        {
            fprintf(stderr, "Error reading line %d from %s\n", i + 1, filename);
            free(buf);
            return nullptr;
        }

        if (buf[0] == '#')
        {
            i--;
            continue;
        }

        buf[strcspn(buf, "\n")] = 0;
        lines[i] = strdup(buf);

        double elapsed = (clock() - lastPrint) * 1000.0 / CLOCKS_PER_SEC;
        if (print_progress && elapsed > 50)
        {
            fprintf(stderr, "\rRead %d lines from %s in %.3f ms\r", i + 1, filename, (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC);
            lastPrint = clock();
        }
    }
    free(buf);
    lines[count] = nullptr;

    fclose(file);
    printf("Read %d definitions from %s in %.3f ms\n", count, filename, (clock() - ctime) * 1000.0 / CLOCKS_PER_SEC);
    return lines;
}

char** grow_string_array(char** array, size_t newSize)
{
    fprintf(stderr, ANSI_CLEAR_LINE ANSI_BLINK "Growing charptr array to %zu" ANSI_RESET "\n", newSize);
    char** newArray = realloc(array, sizeof(char*) * newSize);
    if (newArray == NULL)
    {
        fprintf(stderr, "Memory allocation failed\n");
        exit(EXIT_FAILURE);
    }
    return newArray;
}

int* grow_int_array(int* array, size_t newSize)
{
    fprintf(stderr, ANSI_CLEAR_LINE ANSI_BLINK "Growing int array to %zu" ANSI_RESET "\n", newSize);
    int* newArray = realloc(array, sizeof(int) * newSize);
    if (newArray == NULL)
    {
        fprintf(stderr, "Memory allocation failed\n");
        exit(EXIT_FAILURE);
    }
    return newArray;
}

void ban_ip(char* ip)
{
}

void process_line(char* line, void* cbData)
{
    MatchRules* rules = cbData;

    if (line == NULL)
    {
        fprintf(stderr, "Line is null @ process_line!\n");
        return;
    }


    char* ip = xmalloc(32);
    for (int i = 0; i < 32; i++)
    {
        ip[i] = 0;
    }
    strncpy(ip, line, strcspn(line, " "));

    // printf("meow %s\n", ip);


    if (check_allowed_ip)
        for (int j = 0; rules->allowed_ips[j] != NULL; j++)
        {
            if (streq(ip, rules->allowed_ips[j]))
            {
                if (print_allowed_ip)
                    printf("Allowed IP: \"%s\" (~%s): %s\n", ip, rules->allowed_ips[j], line);
                free(ip);
                return;
            }
        }

    if (check_already_banned_ip)
        for (size_t j = 0; j < rules->banned_ip_count; j++)
        {
            if (strlen(rules->banned_ips[j]) < 2)
            {
                for (size_t k = 0; k < rules->banned_ip_count; k++)
                {
                    printf("rules->banned_ips[%lu] = \"%s\" (#%s) x %d\n", k, rules->banned_ips[k], str_to_hex(rules->banned_ips[k]), rules->banned_ip_hits[k * sizeof(int)]);
                }
                printf("ASSERT: Broken rule #%lu: \"%s\" (#%s)\n", j, rules->banned_ips[j], str_to_hex(rules->banned_ips[j]));
                abort();
            }

            if (streq(ip, rules->banned_ips[j]))
            {
                if (print_already_banned_ip)
                    printf("Banned IP: \"%s\" (rule %lu: %s): %s\n", ip, j, rules->banned_ips[j], line);

                if (strlen(ip) != strlen(rules->banned_ips[j]) || strcmp(ip, rules->banned_ips[j]) != 0)
                {
                    fprintf(stderr, "ASSERT: IP address length did not match! #%lu: \"%s\" (#%s) != \"%s\" (#%s)\n", j, rules->banned_ips[j], str_to_hex(rules->banned_ips[j]), ip,
                            str_to_hex(ip));
                    abort();
                }

                rules->banned_ip_hits[j * sizeof(int)]++;
                fprintf(stderr, "DBG CHK: %s -> %d hits\n", rules->banned_ips[j], rules->banned_ip_hits[j * sizeof(int)]);

                free(ip);
                return;
            }
        }

    if (check_allowed_pattern)
        for (int i = 0; rules->allowed_patterns[i] != NULL; i++)
        {
            if (strstr(line, rules->allowed_patterns[i]) != NULL)
            {
                if (print_allowed_pattern)
                    printf("%15s matched allowed pattern \"%s\": %s\n", ip, rules->allowed_patterns[i], line);
                return;
            }
        }

    for (int i = 0; rules->patterns[i] != NULL; i++)
    {
        if (strstr(line, rules->patterns[i]) != NULL)
        {
            if (rules->banned_ip_count >= rules->banned_ip_size)
            {
                rules->banned_ip_size *= 2;
                rules->banned_ips = grow_string_array(rules->banned_ips, rules->banned_ip_size);
                rules->banned_ip_hits = grow_int_array(rules->banned_ip_hits, rules->banned_ip_size);
            }

            rules->banned_ips[rules->banned_ip_count] = ip;
            rules->banned_ip_hits[rules->banned_ip_count * sizeof(int)] = 1;
            fprintf(stderr, "DBG PTN: %s -> %d hits\n", rules->banned_ips[rules->banned_ip_count], rules->banned_ip_hits[rules->banned_ip_count * sizeof(int)]);
            rules->banned_ip_count++;

            if (print_bans)
                printf("[Bans=%6lu] %15s matched pattern %4d (\"%s\"): %s\n", rules->banned_ip_count, ip, i, rules->patterns[i], line);

            if (print_only) return;

            __pid_t pid = fork();
            if (pid == 0)
            {
                if (execvp("iptables", (char*[]){"iptables", "-A", "INPUT", "-s", ip, "-j", "DROP", nullptr}) != 0)
                {
                    fprintf(stderr, "Error executing iptables: ");
                    perror("");
                }
            }
            else if (pid < 0)
            {
                perror("Fork failed");
                exit(1);
            }
            else
            {
                signal(SIGCHLD, SIG_IGN);
            }

            return;
        }
    }

    if (print_indeterminate)
        printf("[Indeterminate]: %s\n", line);

    free(ip);
}

int main(int argc, char* argv[])
{
    fprintf(stderr, "Hello, World!\n");

    print_allowed_ip = getenv("D_PRINT_ALLOWED_IP") != NULL && strcmp(getenv("D_PRINT_ALLOWED_IP"), "1") == 0;
    print_allowed_pattern = getenv("D_PRINT_ALLOWED_PATTERN") != NULL && strcmp(getenv("D_PRINT_ALLOWED_PATTERN"), "1") == 0;
    print_already_banned_ip = getenv("D_PRINT_BANNED_IP") != NULL && strcmp(getenv("D_PRINT_BANNED_IP"), "1") == 0;
    print_progress = isatty(STDERR_FILENO);

    for (int i = 1; i < argc; i++)
    {
        if (strcmp(argv[i], "--allowed-ips") == 0 && i + 1 < argc)
            allowed_ips_path = argv[++i];
        else if (strcmp(argv[i], "--allowed-patterns") == 0 && i + 1 < argc)
            allowed_patterns_path = argv[++i];
        else if (strcmp(argv[i], "--patterns") == 0 && i + 1 < argc)
            patterns_path = argv[++i];
        else if (strcmp(argv[i], "--access-log") == 0 && i + 1 < argc)
            access_log_path = argv[++i];
        else if (strstr(argv[i], "--print-allowed-ip") != NULL)
            print_allowed_ip = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-allowed-pattern") != NULL)
            print_allowed_pattern = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-already-banned-ip") != NULL)
            print_already_banned_ip = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-bans") != NULL)
            print_bans = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-indeterminate") != NULL)
            print_indeterminate = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-only") != NULL)
            print_only = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--print-progress") != NULL)
            print_progress = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--check-allowed-ip") != NULL)
            check_allowed_ip = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--check-allowed-pattern") != NULL)
            check_allowed_pattern = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--check-already-banned-ip") != NULL)
            check_already_banned_ip = strstr(argv[i], "=false") == NULL;
        else if (strstr(argv[i], "--dump-patterns") != NULL)
            dump_patterns = true;
        else if (strcmp(argv[i], "--debug") == 0)
        {
            print_allowed_ip = true;
            print_allowed_pattern = true;
            print_already_banned_ip = true;
            print_indeterminate = true;
            print_bans = true;
        }
        else if (strstr(argv[i], "--help"))
        {
            printf("Usage: %s [options]\n", argv[0]);
            printf("Options:\n");
            printf("  --allowed-ips <path>                  Path to the allowed IPs file (default: %s)\n", ALLOWED_IPS_FILE);
            printf("  --allowed-patterns <path>             Path to the allowed patterns file (default: %s)\n", ALLOWED_PATTERNS_FILE);
            printf("  --patterns <path>                     Path to the patterns file (default: %s)\n", PATTERNS_FILE);
            printf("  --access-log <path|stdin|->                   Path to the access log file (default: %s)\n", ACCESS_LOG_FILE);
            printf("  --print-allowed-ip[=false]            Print allowed IPs (default: %s)\n", print_allowed_ip ? "true" : "false");
            printf("  --print-allowed-pattern[=false]       Print allowed patterns (default: %s)\n", print_allowed_pattern ? "true" : "false");
            printf("  --print-already-banned-ip[=false]     Print already banned IPs (default: %s)\n", print_already_banned_ip ? "true" : "false");
            printf("  --print-bans[=false]                  Print bans (default: %s)\n", print_bans ? "true" : "false");
            printf("  --print-indeterminate[=false]         Print indeterminate lines (default: %s)\n", print_indeterminate ? "true" : "false");
            printf("  --print-only[=false]                  Print only the lines that match the rules (default: %s)\n", print_only ? "true" : "false");
            printf("  --print-progress[=false]              Print progress while reading file (default: %s)\n", print_progress ? "true" : "false");
            printf("  --debug                               Enable debug mode (prints all information)\n");
            printf("  --check-allowed-ip[=false]            Check allowed IPs (default: %s)\n", check_allowed_ip ? "true" : "false");
            printf("  --check-allowed-pattern[=false]       Check allowed patterns (default: %s)\n", check_allowed_pattern ? "true" : "false");
            printf("  --check-already-banned-ip[=false]     Check already banned IPs (default: %s)\n", check_already_banned_ip ? "true" : "false");
            printf("  --dump-patterns                       Dump pattern table and exit\n");
            return 0;
        }
        else
        {
            fprintf(stderr, "Unknown argument: %s\n", argv[i]);
            return 1;
        }
    }

    if (strcmp(access_log_path, "stdin") == 0 || strcmp(access_log_path, "-") == 0)
    {
        access_log_path = "/dev/stdin";
    }

    print_bool_opt(print_allowed_ip);
    print_bool_opt(print_allowed_pattern);
    print_bool_opt(print_already_banned_ip);
    print_bool_opt(print_indeterminate);
    print_bool_opt(print_bans);
    print_bool_opt(print_only);
    print_bool_opt(print_progress);

    print_bool_opt(check_allowed_ip);
    print_bool_opt(check_allowed_pattern);
    print_bool_opt(check_already_banned_ip);

    print_bool_opt(dump_patterns);

    print_str_opt(allowed_patterns_path);
    print_str_opt(allowed_ips_path);
    print_str_opt(patterns_path);
    print_str_opt(access_log_path);

    MatchRules rules;
    rules.patterns = read_lines(patterns_path);
    rules.allowed_patterns = read_lines(allowed_patterns_path);
    rules.allowed_ips = read_lines(allowed_ips_path);
    rules.banned_ip_size = 8;
    rules.banned_ip_count = 0;
    rules.banned_ips = xmalloc(rules.banned_ip_size * sizeof(char**));
    rules.banned_ip_hits = xmalloc(rules.banned_ip_size * sizeof(int));
    //rules.bannedIps[0] = ;

    if (dump_patterns)
    {
        fprintf(stderr, "Dumping patterns:\n");
        for (int i = 0; rules.patterns[i] != NULL; i++)
        {
            printf("[PATTERN %4d] %s\n", i, rules.patterns[i]);
        }
        for (int i = 0; rules.allowed_patterns[i] != NULL; i++)
        {
            printf("[ALLOWED %4d] %s\n", i, rules.allowed_patterns[i]);
        }
        for (int i = 0; rules.allowed_ips[i] != NULL; i++)
        {
            printf("[ALL. IP %4d] %s\n", i, rules.allowed_ips[i]);
        }
        return 0;
    }

    read_lines_cb(access_log_path, process_line, &rules);

    fprintf(stderr, "Banned %lu IPs:\n", rules.banned_ip_count);

    for (int i = 0; i < rules.banned_ip_count; i++)
    {
        printf("[BAN %4d] %s: %d hits\n", i, rules.banned_ips[i], rules.banned_ip_hits[i * sizeof(int)]);
    }

    for (int i = 0; rules.patterns[i] != NULL; i++)
    {
        free(rules.patterns[i]);
    }
    free(rules.patterns);
    free(rules.allowed_ips);
    for (int i = 0; i < rules.banned_ip_count; i++)
    {
        free(rules.banned_ips[i]);
    }
    free(rules.banned_ips);

    return 0;
}