#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <stdbool.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)
    {
        printf("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;
    }
    printf("[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;
    }
    printf("[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);

    printf("[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){
        free(params);
        zmq_close(requester);
        zmq_ctx_destroy(context);
        return NULL;
    }

    pthread_mutex_lock(&words_mutex);

    len = strlen(reply);
    if(!words_list){
        words_list = (char*)malloc(len + 1);
        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);
    zmq_ctx_destroy(context);

    return NULL;
}

char *map(size_t worker_count, char **worker_ports, char *filename){
    FILE *fp = fopen(filename, "r");
    char *file_content_chunks;
    int worker_index = 0;
    int thread_count = 0;

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

    while (thread_count < worker_count && (file_content_chunks = get_file_content(fp)) != NULL)
    {
        THREAD_ARGS *args = (THREAD_ARGS*)malloc(sizeof(THREAD_ARGS));
        if(!args){
            perror("[MAIN - map] Failed to allocate memory for args");
            free(threads);
            pthread_mutex_destroy(&words_mutex);
            exit(1);
        }
        args->port = atol(worker_ports[worker_index]);
        args->data = strdup(file_content_chunks);
        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;
        }

        free(file_content_chunks);
        worker_index++;
        thread_count++;
    }

    fclose(fp);

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

    free(threads);
    return words_list;
}

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

char *reduce(char *toReduce, 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);
    }

    key_value_pair *head, *tail;

    extract_key_value_pairs(toReduce, &head, &tail);

    size_t thread_count = 0;
    key_value_pair *kv_it = head;

    char *extraction_ptr = toReduce;

    for(int i = 0; kv_it && i < worker_count; i++){
        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[i]);
        args->data = extract_reduce_string(extraction_ptr);
        args->type = RED;

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

        extraction_ptr = toReduce + strlen(args->data) + 1;

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

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

    free(threads);
    return words_list;
}

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
    char *word_list_mapped = map(argc - 2, argv + 2, argv[1]);
    printf("MAP:\n%s\n", word_list_mapped);

    // TODO: REDUCE
    char *word_list_reduced = reduce(word_list_mapped, argc - 2, argv + 2);
    free(word_list_mapped);
    printf("REDUCE:\n%s\n", word_list_reduced);

    // TODO: RIP

    return 0;
}