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

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

#define tolower_str(word) for(int i = 0; i < strlen(word); i++) word[i] = tolower(word[i]);

void* workerFunc(void *args);

char* map(const char *content, key_value_pair **head, key_value_pair **tail);
char* reduce(char *);

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

    printf("Starting...\n");

    int worker_index = 1;
    int worker_count = argc - 1;
    int thread_count = 0;

    pthread_t *threads = (pthread_t*)malloc(sizeof(pthread_t) * worker_count);

    while(thread_count < worker_count){
        int *port = malloc(sizeof(int));
        *port = atoi(argv[worker_index]);
        if(pthread_create(&threads[thread_count++], NULL, workerFunc, port) != 0){
            perror("[MAIN] Failed to create thread");
            free(threads);
            exit(1);
        }
        worker_index++;
        printf("Initialized Thread: %d\n", thread_count);
    }

    printf("Done initialization");

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

    free(threads);

    return 0;
}

TYPE get_request_type(const char *request){
    char type_str[4];
    strncpy(type_str, request, 3);
    type_str[3] = '\0';

    if(strcmp(type_str, "map") == 0) return MAP;
    if(strcmp(type_str, "red") == 0) return RED;
    if(strcmp(type_str, "rip") == 0) return RIP;

    return UNSUPPORTED;
}

void* workerFunc(void *args){
    int port = *(int*)args;
    free(args);
    void *context = zmq_ctx_new();
    if(!context){
        perror("[THREAD] Failed to create context");
        return NULL;
    }
    void *responder = zmq_socket(context, ZMQ_REP);
    if(!responder){
        perror("[THREAD] Failed to create socket");
        zmq_ctx_destroy(context);
        return NULL;
    }

    char addr[22];
    snprintf(addr, sizeof(addr), "tcp://localhost:%d", port);

    int rc = zmq_bind(responder, addr);
    if(rc != 0){
        perror("Failed to connect");
        zmq_close(responder);
        zmq_ctx_destroy(context);
        return NULL;
    }

    while(1){
        key_value_pair *head = NULL;
        key_value_pair *tail = NULL;

        char request[MESSAGE_LENGTH + 1]; // CHECK: off by one error
        int bytes_recv = zmq_recv(responder, request, MESSAGE_LENGTH, 0);

        const TYPE type = get_request_type(request);
        char *content = request + TYPE_LENGTH;

        char *ret;

        switch(type){
            case MAP:
                perror("Handling MAP\n");
                ret = map(content, &head, &tail);
                break;
            case RED:
                perror("Handling RED\n");
                ret = reduce(content);
                break;
            case RIP:
                perror("Handling RIP\n");
            // send errors are redundant to handle because exit
                zmq_send(responder, "rip", TYPE_LENGTH, 0);
                zmq_close(responder);
                zmq_ctx_destroy(context);
                exit(0);
                break;
            case UNSUPPORTED:
                perror("Got unsupported request\n");
            // rip();
                exit(1);
                break;
        }

        if(zmq_send(responder, ret, MESSAGE_LENGTH, 0) != 0){
            perror("Failed to send response");
            zmq_close(responder);
            zmq_ctx_destroy(context);
            exit(1);
        }
    } //loop
}

void add_word(const char *word, key_value_pair **head, key_value_pair **tail){
    char *word_copy = strdup(word);
    tolower_str(word_copy);

    key_value_pair *kv = init_key_value_pair(word, "1");
    if(kv == NULL){
        fprintf(stderr, "Failed to allocate memory for key: %s\n", word);
        free(word_copy);
        return;
    }
    if(add_key_value_pair(kv, head, tail) != 0){
        fprintf(stderr, "Failed to add key value: %s\n", word);
        free(word_copy);
        return;
    }

    free(word_copy);
}

char* get_word(const char *request){
    if(request == NULL || *request == '\0')
        return strdup("");

    int capacity = 16;
    int length = 0;
    char *word = (char*)malloc(capacity);
    if(!word)
        return NULL;

    while(*request && isalpha(*request)){
        if(length + 1 >= capacity){
            capacity *= 2;
            char *new_word = (char*)realloc(word, capacity);
            if(!new_word){
                free(word);
                return NULL;
            }
            word = new_word;
        }
        word[length++] = *request;
        request++;
    }

    word[length] = '\0';
    return word;
}

char* map(const char *content, key_value_pair **head, key_value_pair **tail){
    char *content_iterator = (char*)content;
    void *content_start_address = (void*)content;
    size_t content_length = strlen(content);

    while(content_iterator != NULL && content_iterator < content_start_address + content_length){
        char *word = get_word(content_iterator);
        if(word != NULL){
            if(strcmp(word, "") != 0){
                add_word(word, head, tail);
            }
            content_iterator += strlen(word) + 1;
            free(word);
        }
    }

    key_value_pair *kvp_it = *head;

    char *ret = NULL;

    while(kvp_it != NULL){
        char *kvp2str = kvp_it->toString(kvp_it);
        if(kvp2str == NULL){
            kvp_it = kvp_it->next;
            continue;
        }
        if(!ret){
            ret = strdup(kvp2str);
        }
        else{
            char *tmp = realloc(ret, strlen(ret) + strlen(kvp2str) + 1);
            if(!tmp){
                free(ret);
                perror("Failed to realloc memory for ret");
                free(kvp2str);
                return NULL;
            }
            ret = tmp;
            strcat(ret, kvp2str);
            free(kvp2str);
            kvp_it = kvp_it->next;
        }
    }

    if(!ret){
        perror("Could not extract key value pairs");
        return NULL;
    }

    free_key_value_pairs(head, tail);

    return ret;
}

char* reduce(char *content){
    // free_key_value_pairs();
    return NULL;
}