#include "aoc.h"

typedef struct CellBitMap {
    u64 *bits;
    int size;
} CellBitMap;

static void
bitmap_set_bit(CellBitMap *bitmap, int index) {
    bitmap->bits[index / 64] |= 1ULL << (index % 64);
}

static int
bitmap_popcount(CellBitMap *bitmap) {
    int count = 0;
    for (int i = 0; i < bitmap->size; i++) {
        count += __builtin_popcountll(bitmap->bits[i]);
    }
    return count;
}

static void
bitmap_clear(CellBitMap *bitmap) {
    memset(bitmap->bits, 0, bitmap->size * sizeof(u64));
}

static const int OFFSETS[][2] = {
    { 0, -1},
    { 1,  0},
    { 0,  1},
    {-1,  0},
};

static int 
find_trails_and_compute_rating(Grid *grid, int x, int y, char current_number, CellBitMap *bitmap) {
    ASSERT(grid->grid[y * grid->width + x] == current_number);
    int result = 0;
    if (current_number == '9') {
        int index = y * grid->width + x;
        bitmap_set_bit(bitmap, index);
        result = 1;
    }
    else {
        for (int i = 0; i < 4; i++) {
            int dx = OFFSETS[i][0];
            int dy = OFFSETS[i][1];
            int nx = x + dx;
            int ny = y + dy;
            if (nx >= 0 && nx < grid->width && ny >= 0 && ny < grid->height) {
                char neighbor = grid->grid[ny * grid->width + nx];
                if (neighbor == current_number + 1) {
                    result += find_trails_and_compute_rating(grid, nx, ny, neighbor, bitmap);
                }
            }
        }
    }
    return result;
}

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

    CellBitMap bitmap = {0};
    bitmap.size = (grid.width * grid.height + 63) / 64;
    bitmap.bits = ARENA_ALLOC_ARRAY(arena, u64, bitmap.size);

    i32 part_1 = 0;
    i32 part_2 = 0;
    for (int y = 0; y < grid.height; y++) {
        for (int x = 0; x < grid.width; x++) {
            if (grid.grid[y * grid.width + x] == '0') {
                part_2 += find_trails_and_compute_rating(&grid, x, y, '0', &bitmap);
                int score = bitmap_popcount(&bitmap);
                bitmap_clear(&bitmap);
                part_1 += score;
            }
        }
    }
    printf("%d\n", part_1);
    printf("%d\n", part_2);
}