#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <stdbool.h>

#define QUEUE_CAP 128
#define CHILDREN_CAP 128

#define ARR_SIZE(a) (sizeof(a)/sizeof(*a))

// #define FILENAME "sample.txt"
#define FILENAME "input.txt"

// #define PART1
#define PART2

#define NUMS(X)     \
    X(1, 1, 177622) \
    X(2, 2, 177623) \
    X(3, 3, 177624) \
    X(4, 4, 177625) \
    X(5, 5, 177626) \
    X(6, 6, 177627) \
    X(7, 7, 177628) \
    X(8, 8, 177629) \
    X(9, 9, 177630) \

#define STRS(X)               \
    X(one,   1, 193501607)    \
    X(two,   2, 193507359)    \
    X(three, 3, 210728981597) \
    X(four,  4, 6385231329)   \
    X(five,  5, 6385224815)   \
    X(six,   6, 193505817)    \
    X(seven, 7, 210727692230) \
    X(eight, 8, 210711216854) \
    X(nine,  9, 6385512047)   \

#define TO_STR_ARR(s, v, h) #s,

#define TO_CASE(s, v, h)    \
    case h:                 \
    if(num_str[0] == 0) {   \
        num_str[0] = #v[0]; \
        num_str[1] = #v[0]; \
    } else {                \
        num_str[1] = #v[0]; \
    }                       \
    break;


//  https://www.cse.yorku.ca/~oz/hash.html
unsigned long hash(char *str)
{
    unsigned long hash = 5381;
    int c;

    while ((c = *str++))
        hash = ((hash << 5) + hash) + c;

    return hash;
}

// https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm
typedef struct trie_node {
    struct trie_node *parent;
    struct trie_node *children[CHILDREN_CAP];
    
    struct trie_node *suffix;
    bool end;
} trie;

trie *trie_init(char **dictionary, size_t nwords);
void trie_search(trie *root, char *str, void (*callback)(char *str, size_t len));
void trie_free(trie *root);
void trie_print(trie *root);

size_t sum = 0;
char num_str[3] = {0};

void handle_match(char *str, size_t len)
{      
    char s[len+1]; s[len] = '\0';
    memcpy(s, str, len);
    
    switch(hash(s)) {
        NUMS(TO_CASE)
#ifdef PART2
        STRS(TO_CASE)
#endif
    default:
        fprintf(stderr, "bad hash: %lu (of str %s)", hash(s), s);
    }
}

int main(void)
{
    char *dictionary[] = {
        NUMS(TO_STR_ARR)
#ifdef PART2
        STRS(TO_STR_ARR)
#endif
    };    
    
    trie *root = trie_init(dictionary, ARR_SIZE(dictionary));
    trie_print(root);
    
    FILE *fp = fopen(FILENAME, "r");
    if(!fp) {
        perror("fopen: failed");
        return 1;
    }

    char line[256];
    while(fgets(line, sizeof(line), fp)) {
        trie_search(root, line, handle_match);
        sum += atoi(num_str);
        num_str[0] = 0; num_str[1] = 0;
    }

    fclose(fp);

    printf("Sum: %ld\n", sum);
    
    trie_free(root);
    return 0;
}

static trie *trie_find_suffix(trie *node, size_t idx);
static void trie_set_suffix(trie *node, size_t idx);

static void trie_traverse(trie *parent, void (*callback)(trie *node, size_t idx));
static void trie_add_pattern(trie *node, char *pattern);

trie *trie_init(char **patterns, size_t nwords)
{
    trie *root = malloc(sizeof(trie));
    memset(root, 0, sizeof(trie));
    
    // generate the basic trie
    for(size_t i = 0; i < nwords; i++) {
        trie_add_pattern(root, patterns[i]);
    }

    // calculate suffix links
    trie_traverse(root, trie_set_suffix);

    return root;
}

void trie_search(trie *root, char *str, void (*callback)(char *str, size_t len))
{
    trie *node = root;
    for(size_t i = 0; i < strlen(str)+1; i++) {
        if(node->children[(size_t)str[i]] != NULL) {
            node = node->children[(size_t)str[i]];
            continue;
        }

        if(node->end) {
            int len = 0; trie *tmp = node;
            while(tmp->parent != NULL) {
                tmp = tmp->parent;
                len++;
            }
            callback(&str[i-len], len);
        }
        
        if(node->suffix == NULL) continue;
        
        node = node->suffix;
        i--; continue;
    }
}

void trie_free(trie *root)
{
    for(int i = 0; i < CHILDREN_CAP; i++)
        if(root->children[i] != NULL) {
            trie_free(root->children[i]);
        }
    
    free(root);
}

static void trie_print_child(trie *root, trie *node, size_t depth)
{
    for(int i = 0; i < CHILDREN_CAP; i++) {
        if(node->children[i] != NULL) {
            for(size_t i = 0; i < depth; i++) printf(" ");
            printf("%c (%p) -> (%p) %s\n", i, (void *)node->children[i],
                   (void *)((node->children[i]->suffix == root) ?
                   0 : node->children[i]->suffix),
                   (node->children[i]->end ? "END" : ""));
            trie_print_child(root, node->children[i], depth+2);
        }
    }
}

void trie_print(trie *root)
{
    printf("ROOT (%p)\n", (void *)root);
    trie_print_child(root, root, 2);
}

static trie *trie_find_suffix(trie *node, size_t idx)
{
    if(node->suffix == NULL) return node; // root

    if(node->suffix->children[idx] != NULL) {
        return node->suffix->children[idx];
    } else {
        return trie_find_suffix(node->suffix, idx);
    }
}


static void trie_set_suffix(trie *node, size_t idx)
{
    node->suffix = trie_find_suffix(node->parent, idx);
}

static void trie_traverse(trie *root, void (*callback)(trie *node, size_t idx))
{
    // breadth-first traversal
    
    trie *queue[QUEUE_CAP] = {0};
    size_t queue_start = 0, queue_end = 0;

    queue[queue_end++] = root;

    while(queue_start != queue_end) {
        trie *parent = queue[queue_start++ % QUEUE_CAP];
        for(size_t i = 0; i < CHILDREN_CAP; i++) 
            if(parent->children[i] != NULL) {
                callback(parent->children[i], i);
                queue[queue_end++ % QUEUE_CAP] = parent->children[i];
            }
    }
}

static void trie_add_pattern(trie *node, char *pattern)
{
    if(pattern[0] == '\0') {
        node->end = true;
        return;
    }

    if(node->children[(size_t)pattern[0]] == NULL) {
        trie *child = malloc(sizeof(trie));
        memset(child, 0, sizeof(trie));
        
        child->parent = node;
        node->children[(size_t)pattern[0]] = child;
    }

    trie_add_pattern(node->children[(size_t)pattern[0]], &pattern[1]);
}