#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;
}THREAD_ARGS;

pthread_mutex_t words_mutex;

char *words_list;

char *get_file_content(FILE *fp)
{
    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;
    void *context = zmq_ctx_new();
    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);

    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);

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

    switch(params->type){
        case MAP:
            snprintf(request, len, "map%s", params->data);
            break;
        case RED:
            snprintf(request, len, "red%s", params->data);
            break;
        case RIP:
            snprintf(request, len, "rip");
            break;
        default:
            perror("[THREAD] Got unsupported request");
            free(params);
            free(request);
            zmq_close(requester);
            zmq_ctx_destroy(context);
            return NULL;
    }
    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];
    int bytes_recv = zmq_recv(requester, reply, MESSAGE_LENGTH - 1, 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;
    }

    free(params);

    pthread_mutex_unlock(&words_mutex);

    zmq_close(requester);
    zmq_ctx_destroy(context);

    return NULL;
}

void map(size_t worker_count, char **worker_ports, char *filename) {
    FILE *fp = fopen(filename, "r");
    if (!fp) {
        perror("[MAIN - map] Failed to open file");
        exit(1);
    }

    fseek(fp, 0, SEEK_END);
    size_t file_size = ftell(fp);
    rewind(fp);

    size_t worker_index = 0;
    size_t thread_count = 0;

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

    char buffer[MESSAGE_LENGTH + 1];

    while (!feof(fp)) {
        size_t bytes_read = fread(buffer, sizeof(char), MESSAGE_LENGTH, fp);
        if (bytes_read == 0) break;

        buffer[bytes_read] = '\0';

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

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

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

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

    fclose(fp);

    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(char *toReduce, size_t worker_count, char **worker_ports) {
    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);

        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;

        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){
    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;

        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);

    // MAP
    map(argc - 2, argv + 2, argv[1]);
    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);
    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);

    // print words_list
    // TODO: sort by frequency and alphabetically
    print_words_list(word_list_reduced);
    free(word_list_reduced);

    pthread_mutex_destroy(&words_mutex);
    return 0;
}