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

    fprintf(stderr, "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++;
        fprintf(stderr, "Initialized Thread: %d\n", thread_count);
    }

    fprintf(stderr, "Init DONE\n");

    // 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]; // CHECK: off by one error
        int bytes_recv = zmq_recv(responder, request, MESSAGE_LENGTH - 1, 0);
        if(bytes_recv <= 0){
            perror("Failed to receive request");
            continue;
        }

        request[bytes_recv] = '\0';

        char endpoint[256];
        size_t endpoint_size = sizeof(endpoint);
        if (zmq_getsockopt(responder, ZMQ_LAST_ENDPOINT, endpoint, &endpoint_size) == 0) {
            fprintf(stderr, "Got request from: %s\n", endpoint);
        }

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

        char *ret = NULL;

        switch(type){
            case MAP:
                fprintf(stderr, "Handling MAP\n");
                ret = map(content, &head, &tail);
                break;
            case RED:
                fprintf(stderr, "Handling RED\n");
                ret = reduce(content);
                break;
            case RIP:
                fprintf(stderr, "Handling RIP\n");
            // send errors are redundant to handle because exit
                zmq_send(responder, "rip", TYPE_LENGTH + 1, 0);
                zmq_close(responder);
                zmq_ctx_destroy(context);
                exit(0);
                break;
            case UNSUPPORTED:
                perror("Got unsupported request\n");
            // rip();
                exit(1);
                break;
        }
        fprintf(stderr, "Sending content: %s\n", ret);
        if(zmq_send(responder, ret, strlen(ret), 0) < 0){
            perror("Failed to send response");
            zmq_close(responder);
            zmq_ctx_destroy(context);
            // exit(1);
        }
        free(ret);  // NULL free is safe
        ret = NULL;
    } //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_copy, "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 length = 0;
    char *word = (char*)malloc(sizeof(char));
    if(!word)
        return NULL;

    while(*request && isalpha(*request)){
        char *new_word = (char*)realloc(word, strlen(word) + 2);    // realloc word size + \0 and extra char space
        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;   // skip current read word + one delimiter char
            free(word);
        }
    }

    char *ret = key_value_pair_list_to_string(head);

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

    return ret;
}

char* reduce(char *content){
    key_value_pair *head = NULL;
    key_value_pair *tail = NULL;

    extract_key_value_pairs(content, &head, &tail, true);
    char *ret = key_value_pair_list_to_string(&head);

    free_key_value_pairs(&head, &tail);

    return ret;
}