#include "aoc.h"

typedef struct {
    enum {
        WIRE_CONST,
        WIRE_AND,
        WIRE_OR,
        WIRE_XOR,
    } type;
    union {
        u8 constant;
        struct {
            str a, b;
        } binary;
    };
} Wire;

typedef struct WireEntry {
    struct WireEntry *next;
    u64 hash;
    str key;
    Wire wire;
} WireEntry;

#define WIRE_MAP_EXP 10
#define WIRE_MAP_SIZE (1 << WIRE_MAP_EXP)
#define WIRE_MAP_MASK (WIRE_MAP_SIZE - 1)

_Static_assert((WIRE_MAP_SIZE & WIRE_MAP_MASK) == 0, "WIRE_MAP_SIZE must be a power of 2");

typedef struct WireMap {
    Arena *arena;
    WireEntry *entries[WIRE_MAP_SIZE];
} WireMap;

static void
wm_init(WireMap *map, Arena *arena) {
    memset(map, 0, sizeof(*map));
    map->arena = arena;
}

static u64
hash_str(str s) {
    // FNV-1a hash
    u64 hash = 14695981039346656037ULL;
    for (u64 i = 0; i < s.len; i++) {
        hash ^= s.data[i];
        hash *= 1099511628211ULL;
    }
    return hash;
}

static WireEntry *
wm_get_entry(WireMap *map, str key) {
    u64 hash = hash_str(key);
    u64 index = hash & WIRE_MAP_MASK;
    for (WireEntry *e = map->entries[index]; e; e = e->next) {
        if (e->hash == hash && str_eq(e->key, key)) {
            return e;
        }
    }
    return NULL;
}

static void
wm_put(WireMap *map, str key, Wire wire) {
    WireEntry *entry = NULL;
    u64 hash = hash_str(key);
    u64 index = hash & WIRE_MAP_MASK;
    for (WireEntry *e = map->entries[index]; e; e = e->next) {
        if (e->hash == hash && str_eq(e->key, key)) {
            entry = e;
            break;
        }
    }
    if (!entry) {
        entry = ARENA_ALLOC(map->arena, WireEntry);
        entry->hash = hash;
        entry->key = key;
        entry->next = map->entries[index];
        map->entries[index] = entry;
    }
    entry->wire = wire;
}

static void
debug_dump_wire_map(FILE *f, WireMap *map) {
    for (u64 i = 0; i < WIRE_MAP_SIZE; i++) {
        for (WireEntry *e = map->entries[i]; e; e = e->next) {
            if (e->wire.type == WIRE_CONST) {
                fprintf(f, STR_FMT ": %d\n", STR_ARG(e->key), e->wire.constant);
            } 
            else {
                fprintf(f, STR_FMT ": " STR_FMT "%s" STR_FMT "\n", STR_ARG(e->key), 
                        STR_ARG(e->wire.binary.a), 
                        e->wire.type == WIRE_AND ? " AND " : e->wire.type == WIRE_OR ? " OR " : " XOR ",
                        STR_ARG(e->wire.binary.b));
            }
        }
    }
}

static void
parse_input(Arena temp, WireMap *map, str input) {
    str initial, rules;
    str_split(input, STR("\n\n"), &initial, &rules);
    ASSERT(initial.len > 0);
    ASSERT(rules.len > 0);

    {
        // parse initial
        str line;
        str delim = STR("\n");
        while ((line = str_next_token(&initial, delim)).len > 0) {
            str name, value;
            str_split(line, STR(": "), &name, &value);

            u8 state = (u8) parse_i64(value, temp);
            Wire wire = {
                .type = WIRE_CONST,
                .constant = state,
            };
            wm_put(map, name, wire);
        }
    }

    {
        // parse rules
        str line; 
        str delim = STR("\n");
        while ((line = str_next_token(&rules, delim)).len > 0) {
            line = str_trim(line);
            if (line.len == 0) continue;
            str lhs = str_next_token(&line, STR(" "));
            str op = str_next_token(&line, STR(" "));
            str rhs = str_next_token(&line, STR(" "));
            str arrow = str_next_token(&line, STR("-> "));
            str dest = str_trim(str_next_token(&line, STR("\n")));

            Wire wire;
            if (str_eq(op, STR("AND"))) wire.type = WIRE_AND;
            else if (str_eq(op, STR("OR"))) wire.type = WIRE_OR;
            else if (str_eq(op, STR("XOR"))) wire.type = WIRE_XOR;
            else NOT_REACHABLE();

            wire.binary.a = str_trim(lhs);
            wire.binary.b = str_trim(rhs);
            printf("Inserting wire " STR_FMT "\n", STR_ARG(dest));
            wm_put(map, dest, wire);
        }
    }
}

static u8
evaluate(WireMap *map, str name) {
    WireEntry *entry = wm_get_entry(map, name);
    if (!entry) {
        printf("Wire " STR_FMT " not found\n", STR_ARG(name));
        NOT_REACHABLE();
    }
    ASSERT(entry);
    Wire *wire = &entry->wire;

    if (wire->type == WIRE_CONST) {
        return wire->constant;
    }
    else {
        u8 lhs = evaluate(map, wire->binary.a);
        u8 rhs = evaluate(map, wire->binary.b);

        u8 result = 0;
        switch (wire->type) {
            case WIRE_AND: result = lhs & rhs; break;
            case WIRE_OR: result = lhs | rhs; break;
            case WIRE_XOR: result = lhs ^ rhs; break;
            default: NOT_REACHABLE();
        }
        wire->type = WIRE_CONST;
        wire->constant = result;
        return result;
    }
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(1));
    str input = read_file(arena, argv[1]);

    WireMap map;
    wm_init(&map, arena);
    parse_input(*arena, &map, input);

    u64 part_1 = 0;
    for (u32 index = 0; index < 64; index++) {
        char name[16];
        snprintf(name, sizeof(name), "z%02d", index);
        if (wm_get_entry(&map, cstr(name))) {
            u64 bit = (u64) evaluate(&map, cstr(name)) & 1;
            printf("Bit %d: %lu\n", index, bit);
            part_1 |= (bit << index);
        }
        else {
            break;
        }
    }
    /* debug_dump_wire_map(stderr, &map); */
    printf("%lu\n", part_1);
}