#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <regex.h>
#include <zmq.h>

#include "key_value_pair.h"
#include "transfer_defines.h"

typedef struct{
    int port;
    char *data;
    TYPE type;
    void *context;
}THREAD_ARGS;

pthread_mutex_t words_mutex;

char *words_list;

char *sanitize_input(char *input);

char *get_file_content(FILE *fp)
{
    if (!fp) {
        perror("Failed to open file");
        return NULL;
    }
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    rewind(fp);

    char *content = malloc(file_size + 1);
    if (!content) {
        perror("Failed to allocate memory for entire file");
        fclose(fp);
        return NULL;
    }

    size_t bytes_read = fread(content, 1, file_size, fp);
    content[bytes_read] = '\0';

    fclose(fp);
    return content;
/*
    char *buffer = (char*)malloc(sizeof(char) * MESSAGE_LENGTH);
    if (buffer == NULL)
    {
        perror("Failed to allocate memory for buffer");
        return NULL;
    }

    const size_t bytes_read = fread(buffer, sizeof(char), MESSAGE_LENGTH - 1, fp);
    if (bytes_read == 0)
    {
        perror("File is empty or nothing is left to read");
        free(buffer);
        return NULL;
    }

    buffer[bytes_read] = '\0';

    int delimiter_pos = (int)bytes_read - 1;
    while (delimiter_pos >= 0 && !isalpha(buffer[delimiter_pos])) { delimiter_pos--; }

    if (delimiter_pos < 0)
    {
        fprintf(stderr, "No delimiting character was found -> taking whole buffer");
        delimiter_pos = (int)bytes_read;
    }

    buffer[delimiter_pos + 1] = '\0';

    if (fseek(fp, delimiter_pos - (int) bytes_read + 2, SEEK_CUR) != 0)
    {
        perror("Failed to seek in file");
        free(buffer);
        return NULL;
    }

    return buffer;
    */
}

void *receiverFunc(void *args)
{
    THREAD_ARGS *params = (THREAD_ARGS*)args;
    if(!params){
        perror("[THREAD] invalid arguments");
        return NULL;
    }

    if (params->port <= 0) {
        fprintf(stderr, "[THREAD] Invalid port: %d\n", params->port);
        free(params);
        return NULL;
    }

    void *context = params->context;

    if(!context){
        perror("[THREAD] Failed to create context");
        free(params);
        return NULL;
    }
    void *requester = zmq_socket(context, ZMQ_REQ);
    if(!requester){
        perror("[THREAD] Failed to create socket");
        free(params);
        zmq_ctx_destroy(context);
        return NULL;
    }

    char addr[22];
    snprintf(addr, 21, "tcp://localhost:%d", params->port); // BUG: FIX PORT 0

    int rc = zmq_connect(requester, addr);
    if (rc != 0){
        perror("Failed to connect");
        free(params);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }
    fprintf(stderr, "[THREAD] Connected to worker: %s\n", addr);

    char *data = sanitize_input(params->data);
    if(!data){
        perror("[THREAD] Failed to sanitize input");
        free(params->data);
        free(params);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }

    free(params->data);

    size_t len = strlen(data) + 4;
    char *request = (char*)malloc(sizeof(char) * len);
    if(!request){
        perror("[THREAD] Failed to allocate memory for request");
        free(params);
        free(data);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }


    switch(params->type){
        case MAP:
            snprintf(request, len, "map%s", data);
            break;
        case RED:
            snprintf(request, len, "red%s", data);
            break;
        case RIP:
            snprintf(request, len, "rip");
            break;
        default:
            perror("[THREAD] Got unsupported request");
            free(data);
            free(params);
            free(request);
            zmq_close(requester);
            zmq_ctx_destroy(context);
            return NULL;
    }

    free(data);

    fprintf(stderr, "[THREAD] Send request: %s\n", request);
    if(zmq_send(requester, request, strlen(request), 0) == -1){
        perror("[THREAD] Failed to send request");
        free(params);
        free(request);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }

    free(request);

    fprintf(stderr, "[THREAD]: Waiting for response\n");
    char reply[MESSAGE_LENGTH + 1];
    int bytes_recv = zmq_recv(requester, reply, MESSAGE_LENGTH, 0);
    if (bytes_recv == -1){
        perror("[THREAD] Failed to receive reply");
        free(params);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }

    reply[bytes_recv] = '\0';

    if(strcmp(reply, "rip") == 0){
        perror("[THREAD] Got rip -> shutting down");
        free(params);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }

    pthread_mutex_lock(&words_mutex);

    len = strlen(reply);
    if(!words_list){
        words_list = strdup(reply);
        if(!words_list){
            perror("[THREAD] Failed to allocate memory for words_list");
            free(params);
            pthread_mutex_unlock(&words_mutex);
            zmq_close(requester);
            zmq_ctx_destroy(context);
            return NULL;
        }
    }
    else{
        char *tmp = realloc(words_list , sizeof(char) * (strlen(words_list) + len + 1));
        if(tmp == NULL){
            perror("[THREAD] Failed to realloc memory for words_list");
            free(params);
            pthread_mutex_unlock(&words_mutex);
            zmq_close(requester);
            zmq_ctx_destroy(context);
            return NULL;

        }
        words_list = tmp;
        strcat(words_list, reply);
    }

    free(params);

    pthread_mutex_unlock(&words_mutex);

    zmq_close(requester);

    return NULL;
}

char *sanitize_input(char *input){
    if(!input) return NULL;

    int i = 0;

    char *ret = malloc(sizeof(char));
    if(!ret){
        perror("[MAIN - sanitize_input] Failed to allocate memory for ret");
        exit(1);
    }
    *ret = '\0';

    char *it = input;
    while(!isalpha(*it)) it++;

    bool writeSpace = false;
    while(*it){
        if(isalpha(*it)){
            char *tmp = NULL;
            if(writeSpace){
                tmp = realloc(ret, i + 3);
            }
            else{
                tmp = realloc(ret, i + 2);
            }
            if(!tmp){
                perror("[MAIN - sanitize_input] Failed to realloc memory for ret");
                free(ret);
                return NULL;
            }
            ret = tmp;
            if(writeSpace) ret[i++] = ' ';
            ret[i++] = tolower(*it++);
            writeSpace = false;
        }
        else{
            it++;
            writeSpace = true;
        }
    }
    ret[i] = '\0';
    return ret;
}

void map(size_t worker_count, char **worker_ports, const char *filename, void *context) {
    FILE *fp = fopen(filename, "rb");
    char *raw = get_file_content(fp);
    if(!raw){
        perror("Failed to read file");
        exit(1);
    }
    char *content = sanitize_input(raw);
    free(raw);
    if (!content) {
        pthread_mutex_destroy(&words_mutex);
        exit(1);
    }

    size_t max_threads = worker_count * 2;
    pthread_t *threads = malloc(max_threads * sizeof(pthread_t));
    if (!threads) {
        perror("[MAIN - map] Failed to allocate memory for threads");
        free(content);
        pthread_mutex_destroy(&words_mutex);
        exit(1);
    }

    size_t thread_count = 0;
    size_t worker_index = 0;

    size_t total_len = strlen(content);
    size_t offset = 0;

    while (offset < total_len) {
        size_t chunk_size = (total_len - offset >= PAYLOAD_LENGTH + 1) ? PAYLOAD_LENGTH + 1 : (total_len - offset);

        if (offset + chunk_size < total_len) {
            size_t new_chunk_size = chunk_size;
            while (new_chunk_size > 0) {
                char c = content[offset + new_chunk_size];
                if (!isalpha((unsigned char)c) || (offset + chunk_size < strlen(content) && (content[offset + chunk_size] == ' ' || content[offset + chunk_size] == '\0'))) {
                    break;
                }
                new_chunk_size--;
            }
            if (new_chunk_size == 0) {
                new_chunk_size = chunk_size;
            }
            chunk_size = new_chunk_size;
        }

        if (thread_count == max_threads) {
            max_threads *= 2;
            pthread_t *tmp = realloc(threads, max_threads * sizeof(pthread_t));
            if (!tmp) {
                perror("[MAIN - map] Failed to reallocate memory for threads");
                free(threads);
                free(content);
                pthread_mutex_destroy(&words_mutex);
                exit(1);
            }
            threads = tmp;
        }

        THREAD_ARGS *args = malloc(sizeof(THREAD_ARGS));
        if (!args) {
            perror("[MAIN - map] Failed to allocate memory for args");
            free(threads);
            free(content);
            pthread_mutex_destroy(&words_mutex);
            exit(1);
        }

        args->port = atol(worker_ports[worker_index]);
        args->type = MAP;

        args->data = strndup(content + offset, chunk_size);
        args->context = context;

        if (pthread_create(&threads[thread_count], NULL, receiverFunc, (void*)args) != 0) {
            perror("[MAIN - map] Failed to create thread");
            free(args->data);
            free(args);
            continue;
        }

        offset += chunk_size;
        worker_index = (worker_index + 1) % worker_count;
        thread_count++;
    }
    free(content);

    for (size_t i = 0; i < thread_count; i++) {
        pthread_join(threads[i], NULL);
    }

    free(threads);
}

char *extract_reduce_string(const char * toExtract){
    if(!toExtract) return NULL;
    size_t toExtract_len = strlen(toExtract);
    size_t max_len = toExtract_len < MESSAGE_LENGTH ? toExtract_len : MESSAGE_LENGTH;

    char *end_ptr = (char*)toExtract + max_len - 1;
    while(end_ptr > toExtract && isalpha(*end_ptr)) end_ptr--;

    size_t copy_length = end_ptr - toExtract + 1;

    char *extracted = strndup(toExtract, copy_length);
    if(!extracted){
        perror("[MAIN - reduce - extract-reduce-string] Failed to allocate memory for extracted string");
        return NULL;
    }
    return extracted;
}

void reduce(const char *toReduce, size_t worker_count, char **worker_ports, void *context) {
    if (!toReduce || worker_count == 0) {
        fprintf(stderr, "[MAIN - reduce] Invalid input parameters.\n");
        return;
    }

    pthread_t *threads = malloc(sizeof(pthread_t) * worker_count * 2);
    if (!threads) {
        perror("[MAIN - reduce] Failed to allocate memory for threads");
        pthread_mutex_destroy(&words_mutex);
        exit(1);
    }

    size_t toReduce_len = strlen(toReduce);
    size_t offset = 0;
    size_t worker_index = 0;
    size_t thread_count = 0;

    while (offset < toReduce_len) {
        size_t chunk_size = (toReduce_len - offset >= MESSAGE_LENGTH) ? MESSAGE_LENGTH : (toReduce_len - offset);

        if (offset + chunk_size < toReduce_len) {
            size_t new_size = chunk_size;
            while (new_size > 0) {
                char c = toReduce[offset + new_size];
                if (isdigit((unsigned char)c)) {
                    new_size++;
                    break;
                }
                new_size--;
            }
            if (new_size == 0) {
                new_size = chunk_size;
            }
            chunk_size = new_size;
        }

        THREAD_ARGS *args = (THREAD_ARGS*)malloc(sizeof(THREAD_ARGS));
        if (!args) {
            perror("[MAIN - reduce] Failed to allocate memory for args");
            free(threads);
            pthread_mutex_destroy(&words_mutex);
            exit(1);
        }

        args->port = atol(worker_ports[worker_index]);
        args->data = strndup(toReduce + offset, chunk_size);
        args->type = RED;
        args->context = context;

        if (!args->data) {
            perror("[MAIN - reduce] Failed to allocate memory for chunk data");
            free(args);
            continue;
        }

        if (pthread_create(&threads[thread_count], NULL, receiverFunc, (void*)args) != 0) {
            perror("[MAIN - reduce] Failed to create thread");
            free(args->data);
            free(args);
            continue;
        }

        offset += chunk_size;
        worker_index = (worker_index + 1) % worker_count;
        thread_count++;
    }

    for (size_t i = 0; i < thread_count; i++) {
        pthread_join(threads[i], NULL);
    }

    free(threads);
}


void rip(size_t worker_count, char **worker_ports, void *context){
    pthread_t *threads = malloc(sizeof(pthread_t) * worker_count);
    if(!threads){
        perror("[MAIN- reduce] Failed to allocate memory for threads");
        pthread_mutex_destroy(&words_mutex);
        exit(1);
    }

    size_t thread_count = 0;

    for(int i = 0; i < worker_count; i++){
        THREAD_ARGS *args = (THREAD_ARGS*)malloc(sizeof(THREAD_ARGS));
        if(!args){
            perror("[MAIN - rip] Failed to allocate memory for args");
            free(threads);
            pthread_mutex_destroy(&words_mutex);
            exit(1);
        }
        args->port = atol(worker_ports[i]);
        args->data = strdup("rip");
        args->type = RIP;
        args->context = context;

        if (!args->data) {
            fprintf(stderr, "[MAIN - reduce] Warning: extract_reduce_string() returned NULL\n");
            free(args->data);
            free(args);
            continue;
        }

        if(pthread_create(&threads[thread_count], NULL, receiverFunc, (void*)args) != 0){
            perror("[MAIN - reduce] Failed to create thread");
            free(args);
            continue;
        }
        thread_count++;
    }

    // wait for workers to finish
    for(int i = 0; i < thread_count; i++){
        pthread_join(threads[i], NULL);
    }

    free(threads);
}

int qsort_sort_comparer(const void *kv1, const void *kv2){
    key_value_pair *k1 = *(key_value_pair**)kv1;
    key_value_pair *k2 = *(key_value_pair**)kv2;

    if (!k1->value) return 1;
    if (!k2->value) return -1;

    int freq1 = atol(k1->value);
    int freq2 = atol(k2->value);

    if(freq1 != freq2) return freq2 - freq1;
    return strcmp(k1->key, k2->key);
}

void print_words_list(char *list){
    key_value_pair *head = NULL, *tail = NULL;
    extract_key_value_pairs(list, &head, &tail, false);

    size_t count = 0;
    for (key_value_pair *it = head; it != NULL; it = it->next) {
        count++;
    }

    if (count == 0) {
        printf("word,frequency\n");
        return;
    }

    key_value_pair **kv_array = malloc(count * sizeof(key_value_pair *));
    if (!kv_array) {
        perror("[MAIN - print_words_list] Failed to allocate memory for sorting");
        return;
    }

    size_t index = 0;
    for (key_value_pair *it = head; it != NULL; it = it->next) {
        kv_array[index++] = it;
    }

    qsort(kv_array, count, sizeof(key_value_pair*), qsort_sort_comparer);

    printf("word,frequency\n");
    for (size_t i = 0; i < count; i++) {
        printf("%s,%s\n", kv_array[i]->key, kv_array[i]->value);
    }

    free(kv_array);
    free_key_value_pairs(&head, &tail);
}

int main(const int argc, char **argv) {
    if(argc < 3){
        fprintf(stderr, "Usage: %s <file.txt> <worker port 1> <worker port 2> ... <worker port n>", argv[0]);
        exit(1);
    }

    pthread_mutex_init(&words_mutex, NULL);

    void *context = zmq_ctx_new();
    if(!context){
        perror("[MAIN] zmq_ctx_new failed");
        exit(1);
    }

    // MAP
    map(argc - 2, argv + 2, argv[1], context);
    fprintf(stderr, "MAP:\n%s\n", words_list);

    char *word_list_mapped = strdup(words_list);
    free(words_list);
    words_list = NULL;

    // REDUCE
    reduce(word_list_mapped, argc - 2, argv + 2, context);
    free(word_list_mapped);

    char *word_list_reduced = strdup(words_list);
    free(words_list);
    words_list = NULL;

    fprintf(stderr, "REDUCE:\n%s\n", word_list_reduced);

    // RIP
    rip(argc - 2, argv + 2, context);

    // print words_list
    print_words_list(word_list_reduced);
    free(word_list_reduced);

    pthread_mutex_destroy(&words_mutex);
    zmq_ctx_destroy(context);

    return 0;
}