#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <stdlib.h>
#include <errno.h>

#define TMP1 "tmp1.txt"
#define TMP2 "tmp2.txt"
#define TMP3 "tmp3.txt"
#define TMP4 "tmp4.txt"
#define RESULT "result.txt"

static void write_all(int fd, const char *s)
{
    while (*s != '\0') {
        ssize_t n = write(fd, s, 1);
        if (n < 0) {
            if (errno == EINTR) {
                continue;
            }
            _exit(1);
        }
        s++;
    }
}

static int redirect_input(const char *path)
{
    int fd = open(path, O_RDONLY);

    if (fd < 0) {
        return -1;
    }

    if (dup2(fd, STDIN_FILENO) < 0) {
        close(fd);
        return -1;
    }

    if (close(fd) < 0) {
        return -1;
    }

    return 0;
}

static int redirect_output(const char *path)
{
    int fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0644);

    if (fd < 0) {
        return -1;
    }

    if (dup2(fd, STDOUT_FILENO) < 0) {
        close(fd);
        return -1;
    }

    if (close(fd) < 0) {
        return -1;
    }

    return 0;
}

static int run_stage(const char *name,
                     const char *cmd,
                     char *const args[],
                     const char *input,
                     const char *output)
{
    pid_t pid;
    int status;

    pid = fork();

    if (pid < 0) {
        write_all(STDERR_FILENO, "Error: fork failed\n");
        return -1;
    }

    if (pid == 0) {
        if (input != NULL && redirect_input(input) < 0) {
            write_all(STDERR_FILENO, "Error: input redirection failed\n");
            _exit(127);
        }

        if (output != NULL && redirect_output(output) < 0) {
            write_all(STDERR_FILENO, "Error: output redirection failed\n");
            _exit(127);
        }

        execvp(cmd, args);

        write_all(STDERR_FILENO, "Error: exec failed\n");
        _exit(127);
    }

    write_all(STDOUT_FILENO, "[WordCount] Started ");
    write_all(STDOUT_FILENO, name);
    write_all(STDOUT_FILENO, " process\n");

    while (waitpid(pid, &status, 0) < 0) {
        if (errno == EINTR) {
            continue;
        }

        write_all(STDERR_FILENO, "Error: waitpid failed\n");
        return -1;
    }

    if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
        write_all(STDERR_FILENO, "Error: child process failed\n");
        return -1;
    }

    return 0;
}

static char **build_grep_args(int argc, char *argv[])
{
    int i;
    int pos = 0;
    int hide_names = (argc > 2);
    int extra = hide_names ? 4 : 3;
    char **args = (char **)malloc((argc + extra) * sizeof(char *));

    if (args == NULL) {
        return NULL;
    }

    args[pos++] = "grep";

    if (hide_names) {
        args[pos++] = "-h";
    }

    args[pos++] = "-oE";
    args[pos++] = "[[:alpha:]]+";

    for (i = 1; i < argc; i++) {
        args[pos++] = argv[i];
    }

    args[pos] = NULL;

    return args;
}

int main(int argc, char *argv[])
{
    int ret = 1;
    char **grep_args;
    char *tr_args[] = {"tr", "[:upper:]", "[:lower:]", NULL};
    char *sort_args[] = {"sort", NULL};
    char *uniq_args[] = {"uniq", "-c", NULL};
    char *sort_nr_args[] = {"sort", "-nr", NULL};

    if (argc < 2) {
        write_all(STDERR_FILENO,
                  "Usage: ./WordCount <input_file1> [input_file2 ...]\n");
        return 1;
    }

    grep_args = build_grep_args(argc, argv);

    if (grep_args == NULL) {
        write_all(STDERR_FILENO, "Error: malloc failed\n");
        return 1;
    }

    if (run_stage("grep", "grep", grep_args, NULL, TMP1) != 0) {
        goto cleanup;
    }

    if (run_stage("tr", "tr", tr_args, TMP1, TMP2) != 0) {
        goto cleanup;
    }

    if (run_stage("sort", "sort", sort_args, TMP2, TMP3) != 0) {
        goto cleanup;
    }

    if (run_stage("uniq", "uniq", uniq_args, TMP3, TMP4) != 0) {
        goto cleanup;
    }

    if (run_stage("sort", "sort", sort_nr_args, TMP4, RESULT) != 0) {
        goto cleanup;
    }

    write_all(STDOUT_FILENO, "[WordCount] WordCount completed successfully\n");
    ret = 0;

cleanup:
    free(grep_args);
    return ret;
}
