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

#include "Errors.h"

#define MAX_PATH_LEN 256
#define MAX_REDUCERS 16

static char pgm_help[] =
    "   Sintax: \tWordCountMR_v3 [-m] [-r <nreducers=1>] [-o <output_file>] <input_files>\n"
    "   Example: \tWordCountMR_v3 -m -r 2 -o ./Result/wc_pipe.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_v3] 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) {
        Error("[WordCountMR_v3] fork failed");
    }

    if (pid == 0) {
        execv(child_argv[0], child_argv);
        Error("[WordCountMR_v3] exec failed");
    }

    sprintf(Msg, "[WordCountMR_v3] 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) {
        Error("[WordCountMR_v3] waitpid failed");
    }

    if (!WIFEXITED(status)) {
        Error("[WordCountMR_v3] child did not finish normally");
    }

    if (!status_is_counter && WEXITSTATUS(status) != 0) {
        sprintf(MsgError, "[WordCountMR_v3] %s process failed", stage);
        Error(MsgError);
    }

    return WEXITSTATUS(status);
}

static void close_all_reduce_pipes(int reducer_pipes[][2], int nreducers)
{
    for (int i = 0; i < nreducers; i++) {
        close(reducer_pipes[i][0]);
        close(reducer_pipes[i][1]);
    }
}

static void start_reducers(pid_t reducers[],
                           int reducer_pipes[][2],
                           int nreducers,
                           char *reduce_exec,
                           const char *output_file,
                           int map_pipe[2])
{
    for (int i = 0; i < nreducers; i++) {
        pid_t pid = fork();

        if (pid < 0) {
            Error("[WordCountMR_v3] fork reducer failed");
        }

        if (pid == 0) {
            int fd_out;
            char *args[] = {reduce_exec, NULL};

            close(map_pipe[0]);
            close(map_pipe[1]);

            if (dup2(reducer_pipes[i][0], STDIN_FILENO) < 0) {
                Error("[WordCountMR_v3] dup2 reducer input failed");
            }

            fd_out = open(output_file, O_WRONLY | O_CREAT | O_APPEND, 0644);

            if (fd_out < 0) {
                Error("[WordCountMR_v3] open output failed");
            }

            if (dup2(fd_out, STDOUT_FILENO) < 0) {
                Error("[WordCountMR_v3] dup2 reducer output failed");
            }

            close(fd_out);

            for (int j = 0; j < nreducers; j++) {
                close(reducer_pipes[j][0]);
                close(reducer_pipes[j][1]);
            }

            execv(reduce_exec, args);
            Error("[WordCountMR_v3] exec reducer failed");
        }

        reducers[i] = pid;

        sprintf(Msg, "[WordCountMR_v3] Started WordCountReduce process %d.\n", pid);
        printMessage(Msg, COLOR_GREEN_L);
    }
}

static pid_t start_suffle(int map_pipe[2],
                          int reducer_pipes[][2],
                          int nreducers,
                          char *suffle_exec)
{
    pid_t pid;
    char **args;

    pid = fork();

    if (pid < 0) {
        Error("[WordCountMR_v3] fork suffle failed");
    }

    if (pid == 0) {
        if (dup2(map_pipe[0], STDIN_FILENO) < 0) {
            Error("[WordCountMR_v3] dup2 suffle input failed");
        }

        close(map_pipe[0]);
        close(map_pipe[1]);

        for (int i = 0; i < nreducers; i++) {
            close(reducer_pipes[i][0]);
        }

        args = malloc(sizeof(char *) * (size_t)(nreducers + 3));

        if (args == NULL) {
            Error("[WordCountMR_v3] malloc failed");
        }

        args[0] = suffle_exec;

        args[1] = malloc(16);

        if (args[1] == NULL) {
            Error("[WordCountMR_v3] malloc failed");
        }

        snprintf(args[1], 16, "%d", nreducers);

        for (int i = 0; i < nreducers; i++) {
            args[i + 2] = malloc(16);

            if (args[i + 2] == NULL) {
                Error("[WordCountMR_v3] malloc failed");
            }

            snprintf(args[i + 2], 16, "%d", reducer_pipes[i][1]);
        }

        args[nreducers + 2] = NULL;

        execv(suffle_exec, args);
        Error("[WordCountMR_v3] exec suffle failed");
    }

    sprintf(Msg, "[WordCountMR_v3] Started WordCountSuffle process %d.\n", pid);
    printMessage(Msg, COLOR_GREEN_L);

    return pid;
}

static pid_t start_mapper(char *map_exec,
                          char *input_file,
                          int map_pipe[2],
                          int reducer_pipes[][2],
                          int nreducers)
{
    pid_t pid = fork();

    if (pid < 0) {
        Error("[WordCountMR_v3] fork mapper failed");
    }

    if (pid == 0) {
        char *args[] = {map_exec, input_file, NULL};

        close(map_pipe[0]);

        if (dup2(map_pipe[1], STDOUT_FILENO) < 0) {
            Error("[WordCountMR_v3] dup2 mapper output failed");
        }

        close(map_pipe[1]);

        close_all_reduce_pipes(reducer_pipes, nreducers);

        execv(map_exec, args);
        Error("[WordCountMR_v3] exec mapper failed");
    }

    sprintf(Msg, "[WordCountMR_v3] Started WordCountMap process %d.\n", pid);
    printMessage(Msg, COLOR_GREEN_L);

    return pid;
}

static pid_t start_single_mapper(char *map_exec,
                                 int ninputfiles,
                                 char *input_files[],
                                 int map_pipe[2],
                                 int reducer_pipes[][2],
                                 int nreducers)
{
    pid_t pid = fork();

    if (pid < 0) {
        Error("[WordCountMR_v3] fork mapper failed");
    }

    if (pid == 0) {
        char **args;

        close(map_pipe[0]);

        if (dup2(map_pipe[1], STDOUT_FILENO) < 0) {
            Error("[WordCountMR_v3] dup2 mapper output failed");
        }

        close(map_pipe[1]);

        close_all_reduce_pipes(reducer_pipes, nreducers);

        args = malloc(sizeof(char *) * (size_t)(ninputfiles + 2));

        if (args == NULL) {
            Error("[WordCountMR_v3] malloc failed");
        }

        args[0] = map_exec;

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

        args[ninputfiles + 1] = NULL;

        execv(map_exec, args);
        Error("[WordCountMR_v3] exec mapper failed");
    }

    sprintf(Msg, "[WordCountMR_v3] Started WordCountMap process %d.\n", pid);
    printMessage(Msg, COLOR_GREEN_L);

    return pid;
}

int main(int argc, char *argv[])
{
    bool multiple_mappers = false;
    int par;
    int nreducers = 1;
    int ninputfiles;
    int nmappers;
    int words = 0;
    char *output_file = "./Result/wc_pipe.txt";

    char exec_dir[MAX_PATH_LEN];
    char *map_exec;
    char *suffle_exec;
    char *reduce_exec;

    int map_pipe[2];
    int reducer_pipes[MAX_REDUCERS][2];

    pid_t *mappers;
    pid_t suffle_pid;
    pid_t reducers[MAX_REDUCERS];

    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 (nreducers < 1 || nreducers > MAX_REDUCERS) {
        Error("[WordCountMR_v3] Invalid number of reducers");
    }

    ninputfiles = argc - par;

    if (ninputfiles < 1) {
        Error("[WordCountMR_v3] 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 (pipe(map_pipe) < 0) {
        Error("[WordCountMR_v3] pipe Map-Suffle failed");
    }

    for (int i = 0; i < nreducers; i++) {
        if (pipe(reducer_pipes[i]) < 0) {
            Error("[WordCountMR_v3] pipe Suffle-Reduce failed");
        }
    }

    /*
     * Se crea el fichero final vacío. Después, cada reducer lo abrirá
     * en modo append para escribir su parte del resultado.
     */
    {
        int fd = open(output_file, O_WRONLY | O_CREAT | O_TRUNC, 0644);

        if (fd < 0) {
            Error("[WordCountMR_v3] open final output failed");
        }

        close(fd);
    }

    start_reducers(reducers,
                   reducer_pipes,
                   nreducers,
                   reduce_exec,
                   output_file,
                   map_pipe);

    suffle_pid = start_suffle(map_pipe, reducer_pipes, nreducers, suffle_exec);

    for (int i = 0; i < nreducers; i++) {
        close(reducer_pipes[i][0]);
        close(reducer_pipes[i][1]);
    }

    close(map_pipe[0]);

    mappers = malloc(sizeof(pid_t) * (size_t)nmappers);

    if (mappers == NULL) {
        Error("[WordCountMR_v3] malloc failed");
    }

    if (multiple_mappers) {
        for (int i = 0; i < ninputfiles; i++) {
            mappers[i] = start_mapper(map_exec,
                                      argv[par + i],
                                      map_pipe,
                                      reducer_pipes,
                                      nreducers);
        }
    } else {
        mappers[0] = start_single_mapper(map_exec,
                                         ninputfiles,
                                         &argv[par],
                                         map_pipe,
                                         reducer_pipes,
                                         nreducers);
    }

    close(map_pipe[1]);

    for (int i = 0; i < nmappers; i++) {
        words += wait_child(mappers[i], "WordCountMap", true);
    }

    wait_child(suffle_pid, "WordCountSuffle", false);

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

    free(mappers);
    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_v3] WordCount completed processing %d files and %d words in %.3f secs.\n",
            ninputfiles, words, time);
    printMessage(Msg, COLOR_GREEN_D);

    return 0;
}
