#include "aoc.h"

typedef struct Pattern {
    i8 heights[5];
} Pattern;

typedef DYNAMIC_ARRAY(Pattern) Patterns;

static Pattern
parse_lock_or_key(Grid grid, bool *is_lock) {
    *is_lock = true;
    for (isize x = 0; x < grid.width; x++) {
        if (grid.grid[x] != '#') {
            *is_lock = false;
            break;
        }
    }

    Pattern result = {0};
    ASSERT(grid.width == ARRAY_SIZE(result.heights));
    for (isize x = 0; x < grid.width; x++) {
        i8 height = 0;
        if (*is_lock) {
            for (isize y = 0; y < grid.height && grid.grid[y * grid.width + x] == '#'; y++) {
                height++;
            }
        }
        else {
            for (isize y = grid.height - 1; y >= 0 && grid.grid[y * grid.width + x] == '#'; y--) {
                height++;
            }
        }
        result.heights[x] = height - 1;
    }
    /* printf("Result: %d %d %d %d %d\n", result.heights[0], result.heights[1], result.heights[2], result.heights[3], result.heights[4]); */
    return result;
}

static bool
fits(Pattern lock, Pattern key) {
    for (isize x = 0; x < ARRAY_SIZE(lock.heights); x++) {
        if (lock.heights[x] + key.heights[x] > 5) {
            return false;
        }
    }
    return true;
}

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

    Patterns locks = {0};
    Patterns keys = {0};

    str entry;
    while ((entry = str_next_token(&input, STR("\n\n"))).len > 0) {
        Grid grid = parse_grid(entry, temp);
        bool is_lock;
        Pattern pattern = parse_lock_or_key(grid, &is_lock);
        if (is_lock) {
            *push(&locks, arena) = pattern;
        }
        else {
            *push(&keys, arena) = pattern;
        }
        arena_clear(temp);
    }

    i32 part_1 = 0;
    for (isize lock_index = 0; lock_index < locks.len; lock_index++) {
        Pattern lock = locks.data[lock_index];
        for (isize key_index = 0; key_index < keys.len; key_index++) {
            Pattern key = keys.data[key_index];
            if (fits(lock, key)) {
                part_1++;
                /* printf("Lock %ld fits key %ld\n", lock_index, key_index); */
            }
        }
    }
    printf("%d\n", part_1);
}