#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <time.h>
#include <locale.h>
#include <stdbool.h>

#include "Suffle.h"
#include "Reduce.h"
#include "Errors.h"

#define MAX_PATH_LEN 256

static char pgm_help[] =
    "   Sintax: \tWordCountMR_v2b [-m] [-r <nreducers=1>] [-o <output_file>] <input_files>\n"
    "   Example: \tWordCountMR_v2b -m -r 2 -o ./Result/wc.txt ./Input/README.txt ./Input/LEEME.txt\n";

static void get_exec_dir(const char *argv0, char *dir, size_t size)
{
    const char *last_slash = strrchr(argv0, '/');
    size_t len;

    if (last_slash == NULL) {
        snprintf(dir, size, ".");
        return;
    }

    len = (size_t)(last_slash - argv0);

    if (len >= size) {
        len = size - 1;
    }

    memcpy(dir, argv0, len);
    dir[len] = '\0';
}

static char *build_exec_path(const char *dir, const char *program)
{
    size_t len_dir = strlen(dir);
    size_t len_program = strlen(program);
    char *path = malloc(len_dir + 1 + len_program + 1);

    if (path == NULL) {
        Error("[WordCountMR_v2b] malloc failed");
    }

    memcpy(path, dir, len_dir);
    path[len_dir] = '/';
    memcpy(path + len_dir + 1, program, len_program);
    path[len_dir + 1 + len_program] = '\0';

    return path;
}

static pid_t start_child(char *const child_argv[], const char *stage)
{
    pid_t pid = fork();

    if (pid < 0) {
        sprintf(MsgError, "[WordCountMR_v2b] Error creating process");
        Error(MsgError);
    }

    if (pid == 0) {
        execv(child_argv[0], child_argv);

        sprintf(MsgError, "[WordCountMR_v2b] Error executing process");
        Error(MsgError);
    }

    sprintf(Msg, "[WordCountMR_v2b] Started %s process %d.\n", stage, pid);
    printMessage(Msg, COLOR_GREEN_L);

    return pid;
}

static int wait_child(pid_t pid, const char *stage, bool status_is_counter)
{
    int status;

    if (waitpid(pid, &status, 0) < 0) {
        sprintf(MsgError, "[WordCountMR_v2b] Error waiting for process");
        Error(MsgError);
    }

    if (!WIFEXITED(status)) {
        sprintf(MsgError, "[WordCountMR_v2b] Process did not finish normally");
        Error(MsgError);
    }

    /*
     * En WordCountMap, el código de salida se utiliza para devolver el número
     * de palabras procesadas. Por eso no se interpreta como error que sea
     * distinto de cero.
     */
    if (!status_is_counter && WEXITSTATUS(status) != 0) {
        sprintf(MsgError, "[WordCountMR_v2b] %s process failed", stage);
        Error(MsgError);
    }

    return WEXITSTATUS(status);
}

static int run_single_map(char *map_exec, int ninputfiles, char *input_files[])
{
    char **map_argv;
    pid_t pid;
    int words;

    map_argv = malloc(sizeof(char *) * (ninputfiles + 2));

    if (map_argv == NULL) {
        Error("[WordCountMR_v2b] malloc failed");
    }

    map_argv[0] = map_exec;

    for (int i = 0; i < ninputfiles; i++) {
        map_argv[i + 1] = input_files[i];
    }

    map_argv[ninputfiles + 1] = NULL;

    pid = start_child(map_argv, "WordCountMap");
    words = wait_child(pid, "WordCountMap", true);

    free(map_argv);

    return words;
}

static int run_multiple_maps(char *map_exec, int ninputfiles, char *input_files[])
{
    pid_t *pids;
    char ***args;
    char **ids;
    int totalwords = 0;

    pids = malloc(sizeof(pid_t) * ninputfiles);
    args = malloc(sizeof(char **) * ninputfiles);
    ids = malloc(sizeof(char *) * ninputfiles);

    if (pids == NULL || args == NULL || ids == NULL) {
        Error("[WordCountMR_v2b] malloc failed");
    }

    for (int i = 0; i < ninputfiles; i++) {
        ids[i] = malloc(16);
        args[i] = malloc(sizeof(char *) * 5);

        if (ids[i] == NULL || args[i] == NULL) {
            Error("[WordCountMR_v2b] malloc failed");
        }

        snprintf(ids[i], 16, "%d", i);

        args[i][0] = map_exec;
        args[i][1] = "-i";
        args[i][2] = ids[i];
        args[i][3] = input_files[i];
        args[i][4] = NULL;

        pids[i] = start_child(args[i], "WordCountMap");
    }

    for (int i = 0; i < ninputfiles; i++) {
        totalwords += wait_child(pids[i], "WordCountMap", true);
        free(ids[i]);
        free(args[i]);
    }

    free(ids);
    free(args);
    free(pids);

    return totalwords;
}

static void run_suffle(char *suffle_exec, int nmappers, int nreducers)
{
    char nmappers_str[16];
    char nreducers_str[16];
    char *suffle_argv[4];
    pid_t pid;

    snprintf(nmappers_str, sizeof(nmappers_str), "%d", nmappers);
    snprintf(nreducers_str, sizeof(nreducers_str), "%d", nreducers);

    suffle_argv[0] = suffle_exec;
    suffle_argv[1] = nmappers_str;
    suffle_argv[2] = nreducers_str;
    suffle_argv[3] = NULL;

    pid = start_child(suffle_argv, "WordCountSuffle");
    wait_child(pid, "WordCountSuffle", false);
}

static void run_reducers(char *reduce_exec, char *output_file, int nreducers)
{
    pid_t *pids;
    char ***args;
    char **ids;
    char nreducers_str[16];

    pids = malloc(sizeof(pid_t) * nreducers);
    args = malloc(sizeof(char **) * nreducers);
    ids = malloc(sizeof(char *) * nreducers);

    if (pids == NULL || args == NULL || ids == NULL) {
        Error("[WordCountMR_v2b] malloc failed");
    }

    snprintf(nreducers_str, sizeof(nreducers_str), "%d", nreducers);

    for (int i = 0; i < nreducers; i++) {
        ids[i] = malloc(16);
        args[i] = malloc(sizeof(char *) * 5);

        if (ids[i] == NULL || args[i] == NULL) {
            Error("[WordCountMR_v2b] malloc failed");
        }

        snprintf(ids[i], 16, "%d", i);

        args[i][0] = reduce_exec;
        args[i][1] = output_file;
        args[i][2] = nreducers_str;
        args[i][3] = ids[i];
        args[i][4] = NULL;

        pids[i] = start_child(args[i], "WordCountReduce");
    }

    for (int i = 0; i < nreducers; i++) {
        wait_child(pids[i], "WordCountReduce", false);
        free(ids[i]);
        free(args[i]);
    }

    free(ids);
    free(args);
    free(pids);
}

int main(int argc, char *argv[])
{
    bool multiple_mappers = false;
    int par;
    int nreducers = 1;
    int ninputfiles;
    int nmappers;
    int words;
    char *output_file = NULL;
    char exec_dir[MAX_PATH_LEN];
    char *map_exec;
    char *suffle_exec;
    char *reduce_exec;
    struct timespec start, finish;
    float time;

    if (argc < 2) {
        printMessage("Error number of arguments.\n", COLOR_RED);
        printMessage(pgm_help, COLOR_RED);
        exit(1);
    }

    for (par = 1; par < argc && argv[par][0] == '-'; par++) {
        switch (argv[par][1]) {
            case 'm':
                multiple_mappers = true;
                break;

            case 'r':
                if (par + 1 >= argc) {
                    printMessage("Error missing value for -r.\n", COLOR_RED);
                    printMessage(pgm_help, COLOR_RED);
                    exit(1);
                }

                nreducers = atoi(argv[++par]);
                break;

            case 'o':
                if (par + 1 >= argc) {
                    printMessage("Error missing value for -o.\n", COLOR_RED);
                    printMessage(pgm_help, COLOR_RED);
                    exit(1);
                }

                output_file = argv[++par];
                break;

            default:
                printMessage("Error optional argument is incorrect.\n", COLOR_RED);
                printMessage(pgm_help, COLOR_RED);
                exit(1);
        }
    }

    if (output_file == NULL) {
        output_file = DReduceOutputFileName;
    }

    if (nreducers < 1 || nreducers > DMAX_REDUCERS) {
        Error("[WordCountMR_v2b] Invalid number of reducers");
    }

    ninputfiles = argc - par;

    if (ninputfiles < 1) {
        Error("[WordCountMR_v2b] No input files provided");
    }

    if (ninputfiles == 1) {
        multiple_mappers = false;
    }

    nmappers = multiple_mappers ? ninputfiles : 1;

    setlocale(LC_ALL, "");
    clock_gettime(CLOCK_MONOTONIC, &start);

    get_exec_dir(argv[0], exec_dir, sizeof(exec_dir));

    map_exec = build_exec_path(exec_dir, "WordCountMap");
    suffle_exec = build_exec_path(exec_dir, "WordCountSuffle");
    reduce_exec = build_exec_path(exec_dir, "WordCountReduce");

    if (multiple_mappers) {
        words = run_multiple_maps(map_exec, ninputfiles, &argv[par]);
    } else {
        words = run_single_map(map_exec, ninputfiles, &argv[par]);
    }

    run_suffle(suffle_exec, nmappers, nreducers);
    run_reducers(reduce_exec, output_file, nreducers);

    free(map_exec);
    free(suffle_exec);
    free(reduce_exec);

    clock_gettime(CLOCK_MONOTONIC, &finish);

    time = (finish.tv_sec - start.tv_sec) +
           ((finish.tv_nsec - start.tv_nsec) / 1000000000.0);

    sprintf(Msg,
            "[WordCountMR_v2b] WordCount completed processing %d files and %d words in %.3f secs.\n",
            ninputfiles, words, time);
    printMessage(Msg, COLOR_GREEN_D);

    return 0;
}
