#include "aoc.h"
#include <stdlib.h>

static inline bool
is_even(i64 x) { return (x & 1) == 0; }

static inline i64
count_digits(i64 x) {
    ASSERT(x >= 0);
    if (x < 10ull) return 1;
    if (x < 100ull) return 2;
    if (x < 1000ull) return 3;
    if (x < 10000ull) return 4;
    if (x < 100000ull) return 5;
    if (x < 1000000ull) return 6;
    if (x < 10000000ull) return 7;
    if (x < 100000000ull) return 8;
    if (x < 1000000000ull) return 9;
    if (x < 10000000000ull) return 10;
    if (x < 100000000000ull) return 11;
    if (x < 1000000000000ull) return 12;
    if (x < 10000000000000ull) return 13;
    if (x < 100000000000000ull) return 14;
    if (x < 1000000000000000ull) return 15;
    if (x < 10000000000000000ull) return 16;
    if (x < 100000000000000000ull) return 17;
    if (x < 1000000000000000000ull) return 18;
    if (x < 10000000000000000000ull) return 19;
    return 20;
}

static const i64
POWERS_OF_TEN[] = {
    1ull,
    10ull,
    100ull,
    1000ull,
    10000ull,
    100000ull,
    1000000ull,
    10000000ull,
    100000000ull,
    1000000000ull,
    10000000000ull,
    100000000000ull,
    1000000000000ull,
    10000000000000ull,
    100000000000000ull,
    1000000000000000ull,
    10000000000000000ull,
    100000000000000000ull,
    1000000000000000000ull,
    10000000000000000000ull,
};

static u64
hash_stone_iterations(i64 stone, i32 iterations) {
    return ((u64) stone * 73856093ull) ^ ((u64) iterations * 2654435761ull);
}

#define CACHE_SIZE (1 << 18) // empirically determined
_Static_assert((CACHE_SIZE & (CACHE_SIZE - 1)) == 0, "CACHE_SIZE must be a power of 2");
#define CACHE_MASK (CACHE_SIZE - 1)

typedef struct CacheSlot {
    struct CacheSlot *next;
    i64 stone;
    i32 iterations;
    i64 result;
} CacheSlot;


typedef struct Cache {
    CacheSlot *slots[CACHE_SIZE];
    Arena *arena;
} Cache;

static void
cache_init(Cache *cache, Arena *arena) {
    cache->arena = arena;
}

static bool
cache_get(Cache *cache, i64 stone, i32 iterations, i64 *result) {
    u64 hash = hash_stone_iterations(stone, iterations);
    u64 index = hash & CACHE_MASK;
    for (CacheSlot *slot = cache->slots[index]; slot; slot = slot->next) {
        if (slot->stone == stone && slot->iterations == iterations) {
            *result = slot->result;
            return true;
        }
    }
    return false;
}

static void
cache_put(Cache *cache, i64 stone, i32 iterations, i64 result) {
    u64 hash = hash_stone_iterations(stone, iterations);
    u64 index = hash & CACHE_MASK;
    for (CacheSlot *slot = cache->slots[index]; slot; slot = slot->next) {
        if (slot->stone == stone && slot->iterations == iterations) {
            slot->result = result;
            return;
        }
    }
    // not found
    CacheSlot *slot = ARENA_ALLOC(cache->arena, CacheSlot);
    slot->stone = stone;
    slot->iterations = iterations;
    slot->result = result;
    slot->next = cache->slots[index];
    cache->slots[index] = slot;
}

static i64
evolve(Cache *cache, i64 stone, i32 iterations) {
    if (iterations == 0) {
        return 1;
    }
    else {
        i64 result = 0;
        if (cache_get(cache, stone, iterations, &result)) {
            return result;
        }
        else if (stone == 0) {
            result = evolve(cache, 1, iterations - 1);
        }
        else {
            i64 digit_count = count_digits(stone);
            if (is_even(digit_count)) {
                i64 pot = POWERS_OF_TEN[digit_count / 2];
                i64 left = stone / pot;
                i64 right = stone % pot;
                result = evolve(cache, left, iterations - 1) + evolve(cache, right, iterations - 1);
            }
            else {
                i64 new_stone = stone * 2024;
                result = evolve(cache, new_stone, iterations - 1);
            }
        }
        cache_put(cache, stone, iterations, result);
        return result;
    }
}

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

    Cache *cache = ARENA_ALLOC(arena, Cache);
    cache_init(cache, arena);
    
    str iter = input;
    str token;
    i64 part_1 = 0;
    i64 part_2 = 0;
    while ((token = str_next_token(&iter, STR(" "))).len > 0) {
        i64 stone = parse_i64(token, *arena);
        part_1 += evolve(cache, stone, 25);
        part_2 += evolve(cache, stone, 75);
    }
    printf("%ld\n", part_1);
    printf("%ld\n", part_2);
}