#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <locale.h>

#include "Errors.h"

#define INITIAL_CAPACITY 1024
#define READ_BUFFER_SIZE 4096

typedef struct {
    char *word;
    char *line;
} Tuple;

static void *xmalloc(size_t size)
{
    void *ptr = malloc(size);

    if (ptr == NULL) {
        Error("[WordCountSuffle] malloc failed");
    }

    return ptr;
}

static char *xstrndup(const char *s, size_t n)
{
    char *copy = xmalloc(n + 1);

    memcpy(copy, s, n);
    copy[n] = '\0';

    return copy;
}

static void write_all(int fd, const char *buffer, size_t length)
{
    size_t written = 0;

    while (written < length) {
        ssize_t n = write(fd, buffer + written, length - written);

        if (n < 0) {
            Error("[WordCountSuffle] write failed");
        }

        written += (size_t)n;
    }
}

static unsigned long hash_word(const char *word)
{
    unsigned long hash = 5381;
    int c;

    while ((c = *word++) != '\0') {
        hash = ((hash << 5) + hash) + (unsigned long)c;
    }

    return hash;
}

static int compare_tuples(const void *a, const void *b)
{
    const Tuple *ta = (const Tuple *)a;
    const Tuple *tb = (const Tuple *)b;

    return strcmp(ta->word, tb->word);
}

static void add_tuple(Tuple **tuples,
                      int *count,
                      int *capacity,
                      const char *line,
                      size_t line_len)
{
    size_t i = 0;
    size_t word_len;
    char *word;
    char *tuple_line;

    while (i < line_len && line[i] != ' ' && line[i] != '\t') {
        i++;
    }

    word_len = i;

    if (word_len == 0) {
        return;
    }

    if (*count >= *capacity) {
        *capacity *= 2;
        *tuples = realloc(*tuples, sizeof(Tuple) * (size_t)(*capacity));

        if (*tuples == NULL) {
            Error("[WordCountSuffle] realloc failed");
        }
    }

    word = xstrndup(line, word_len);
    tuple_line = xstrndup(line, line_len);

    (*tuples)[*count].word = word;
    (*tuples)[*count].line = tuple_line;
    (*count)++;
}

static void read_all_tuples(Tuple **tuples, int *count, int *capacity)
{
    char read_buffer[READ_BUFFER_SIZE];
    char *line_buffer = NULL;
    size_t line_capacity = 0;
    size_t line_length = 0;

    for (;;) {
        ssize_t n = read(STDIN_FILENO, read_buffer, sizeof(read_buffer));

        if (n < 0) {
            Error("[WordCountSuffle] read failed");
        }

        if (n == 0) {
            break;
        }

        for (ssize_t i = 0; i < n; i++) {
            char c = read_buffer[i];

            if (c == '\n') {
                if (line_length > 0) {
                    add_tuple(tuples, count, capacity, line_buffer, line_length);
                    line_length = 0;
                }
            } else {
                if (line_length + 1 >= line_capacity) {
                    line_capacity = (line_capacity == 0) ? 128 : line_capacity * 2;
                    line_buffer = realloc(line_buffer, line_capacity);

                    if (line_buffer == NULL) {
                        Error("[WordCountSuffle] realloc failed");
                    }
                }

                line_buffer[line_length++] = c;
            }
        }
    }

    if (line_length > 0) {
        add_tuple(tuples, count, capacity, line_buffer, line_length);
    }

    free(line_buffer);
}

int main(int argc, char *argv[])
{
    int nreducers;
    int *reduce_fds;
    Tuple *tuples;
    int count = 0;
    int capacity = INITIAL_CAPACITY;

    if (argc < 3) {
        Error("[WordCountSuffle] Usage: WordCountSuffle <nreducers> <fd1> ... <fdN>");
    }

    nreducers = atoi(argv[1]);

    if (nreducers < 1 || argc != nreducers + 2) {
        Error("[WordCountSuffle] Invalid number of reducers");
    }

    setlocale(LC_ALL, "");

    reduce_fds = xmalloc(sizeof(int) * (size_t)nreducers);
    tuples = xmalloc(sizeof(Tuple) * (size_t)capacity);

    for (int i = 0; i < nreducers; i++) {
        reduce_fds[i] = atoi(argv[i + 2]);
    }

    read_all_tuples(&tuples, &count, &capacity);

    qsort(tuples, (size_t)count, sizeof(Tuple), compare_tuples);

    for (int i = 0; i < count; i++) {
        int reducer = (int)(hash_word(tuples[i].word) % (unsigned long)nreducers);

        write_all(reduce_fds[reducer], tuples[i].line, strlen(tuples[i].line));
        write_all(reduce_fds[reducer], "\n", 1);
    }

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

    for (int i = 0; i < count; i++) {
        free(tuples[i].word);
        free(tuples[i].line);
    }

    free(tuples);
    free(reduce_fds);

    return 0;
}
