#include "aoc.h"

static bool
can_push(Grid *grid, i32 x, i32 y, i32 dx, i32 dy) {
    i32 index = y * grid->width + x;
    i32 next_index = (y + dy) * grid->width + (x + dx);
    u8 cell = grid->grid[index];
    u8 next_cell = grid->grid[index];
    switch (next_cell) {
        case '.': return true;
        case '#': return false;
        case '@':
        case 'O': return can_push(grid, x + dx, y + dy, dx, dy);
        case '[':
            if (dy == 0) return can_push(grid, x + dx, y, dx, dy);
            else         return can_push(grid, x, y + dy, dx, dy) && can_push(grid, x + 1, y + dy, dx, dy);
        case ']':
            if (dy == 0) return can_push(grid, x + dx, y, dx, dy);
            else         return can_push(grid, x, y + dy, dx, dy) && can_push(grid, x - 1, y + dy, dx, dy);
        default: NOT_REACHABLE();
    }
}

static bool
push_block(Grid *grid, i32 x, i32 y, i32 dx, i32 dy) {
    if (!can_push(grid, x, y, dx, dy)) return false;

    u8 cell = grid_at(grid, x, y);
    if (cell == '.') {
        return true;
    }
    else if (cell == '#') {
        return false;
    }
    else if (cell == 'O' || cell == '@') {
        i32 nx = x + dx;
        i32 ny = y + dy;
        if (push_block(grid, nx, ny, dx, dy)) {
            ASSERT(can_push(grid, x, y, dx, dy));
            grid->grid[ny * grid->width + nx] = grid->grid[y * grid->width + x];
            grid->grid[y * grid->width + x] = '.';
            return true;
        }
        else {
            return false;
        }
    }
    else if (cell == '[' || cell == ']') {
        i32 nx = x + dx;
        i32 ny = y + dy;
        if (dy == 0) {
            if (push_block(grid, nx, ny, dx, dy)) {
                ASSERT(can_push(grid, x, y, dx, dy));
                grid->grid[ny * grid->width + nx] = grid->grid[y * grid->width + x];
                grid->grid[y * grid->width + x] = '.';
                return true;
            }
            else {
                return false;
            }
        }
        else {
            ASSERT(dy != 0);
            i32 partner_x = (cell == '[') ? nx + 1 : nx - 1;
            if (push_block(grid, nx, ny, dx, dy) 
                && push_block(grid, partner_x, ny, dx, dy)) 
            {
                ASSERT(can_push(grid, x, y, dx, dy));
                grid->grid[ny * grid->width + nx] = grid->grid[y * grid->width + x];
                grid->grid[ny * grid->width + partner_x] = grid->grid[y * grid->width + partner_x];
                grid->grid[y * grid->width + x] = '.';
                grid->grid[y * grid->width + partner_x] = '.';
                return true;
            }
            else {
                return false;
            }
        }
    }
    else NOT_REACHABLE();
}

static void
print_grid(Grid *grid) {
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            putchar(grid_at(grid, x, y));
        }
        putchar('\n');
    }
}

static Grid
widen(Arena *arena, Grid *grid) {
    Grid result = {
        .width = 2 * grid->width,
        .height = grid->height,
        .grid = ARENA_ALLOC_ARRAY(arena, u8, 2 * grid->width * grid->height)
    };
    u8 *ptr = result.grid;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            u8 cell = grid_at(grid, x, y);
            if (cell == '.') {
                *ptr++ = '.';
                *ptr++ = '.';
            }
            else if (cell == '#') {
                *ptr++ = '#';
                *ptr++ = '#';
            }
            else if (cell == 'O') {
                *ptr++ = '[';
                *ptr++ = ']';
            }
            else if (cell == '@') {
                *ptr++ = '@';
                *ptr++ = '.';
            }
            else {
                NOT_REACHABLE();
            }
        }
    }
    ASSERT(ptr == result.grid + result.width * result.height);
    return result;
}

static i32
compute_total(Grid *grid) {
    i32 result = 0;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            u8 cell = grid_at(grid, x, y);
            if (cell == 'O' || cell == '[') {
                result += 100 * y + x;
            }
        }
    }
    return result;
}

static void
find_robot(Grid *grid, i32 *x, i32 *y) {
    i32 robot_x = -1, robot_y = -1;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            if (grid_at(grid, x, y) == '@') {
                robot_x = x;
                robot_y = y;
            }
        }
    }
    ASSERT(robot_x != -1 && robot_y != -1);
    *x = robot_x;
    *y = robot_y;
}

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

    str grid_input, instructions_str;
    str_split(input, STR("\n\n"), &grid_input, &instructions_str);

    Grid original = parse_grid(grid_input, arena);
    Grid grid = parse_grid(grid_input, arena);

    i32 robot_x, robot_y;
    find_robot(&grid, &robot_x, &robot_y);

    u8 *instructions = ARENA_ALLOC_ARRAY(arena, u8, instructions_str.len);
    isize instruction_count = 0;
    for (isize i = 0; i < instructions_str.len; i++) {
        u8 ch = instructions_str.data[i];
        if (ch != '\n') {
            instructions[instruction_count++] = ch;
        }
    }

    // Part 1
    for (isize i = 0; i < instruction_count; i++) {
        u8 ch = instructions[i];
        i32 dx = 0, dy = 0;
        if (ch == '<') dx = -1;
        else if (ch == '>') dx = 1;
        else if (ch == '^') dy = -1;
        else if (ch == 'v') dy = 1;
        else NOT_REACHABLE();

        if (push_block(&grid, robot_x, robot_y, dx, dy)) {
            robot_x += dx;
            robot_y += dy;
        }
    }

    i32 part_1 = compute_total(&grid);
    printf("%d\n", part_1);

    // Part 2
    Grid wide_grid = widen(arena, &original);
    find_robot(&wide_grid, &robot_x, &robot_y);

    for (isize i = 0; i < instruction_count; i++) {
        u8 ch = instructions[i];
        i32 dx = 0, dy = 0;
        if (ch == '<') dx = -1;
        else if (ch == '>') dx = 1;
        else if (ch == '^') dy = -1;
        else if (ch == 'v') dy = 1;
        else NOT_REACHABLE();

        if (push_block(&wide_grid, robot_x, robot_y, dx, dy)) {
            robot_x += dx;
            robot_y += dy;
        }

    }
    i32 part_2 = compute_total(&wide_grid);
    printf("%d\n", part_2);
}