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

static const i32 DX[4] = {0, 1, 0, -1};
static const i32 DY[4] = {-1, 0, 1, 0};

typedef struct {
    i32 start_x, start_y;
    i32 end_x, end_y;
} Cheat;

typedef struct {
    i32 *dist;
    i32 width, height;
} DistanceMap;

typedef struct {
    i32 x, y;
} Vec2i;

static Vec2i
vec2i(i32 x, i32 y) {
    Vec2i result = { .x = x, .y = y };
    return result;
}

typedef struct {
    Vec2i *points;
    i32 write, read;
    isize size;
} Queue;

static inline i32
queue_next(Queue *queue, i32 index) {
    return (index + 1) % queue->size;
}

static inline void
queue_push(Queue *queue, Vec2i pos) {
    ASSERT(queue_next(queue, queue->write) != queue->read);
    queue->points[queue->write] = pos;
    queue->write = queue_next(queue, queue->write);
}

static inline Vec2i
queue_pop(Queue *queue) {
    ASSERT(queue->read != queue->write);
    Vec2i result = queue->points[queue->read];
    queue->read = queue_next(queue, queue->read);
    return result;
}

static inline bool
queue_empty(Queue *queue) {
    return queue->write == queue->read;
}

static DistanceMap
compute_distances(Arena *arena, Grid grid, i32 x, i32 y) {
    DistanceMap map = {
        .width = grid.width,
        .height = grid.height,
        .dist = ARENA_ALLOC_ARRAY(arena, i32, grid.width * grid.height)
    };
    for (i32 i = 0; i < grid.width * grid.height; i++) {
        map.dist[i] = INT32_MAX;
    }
    map.dist[y * map.width + x] = 0;
    Queue queue = {.size = 1024, .points = ARENA_ALLOC_ARRAY(arena, Vec2i, 1024) };
    queue_push(&queue, vec2i(x, y));
    while (!queue_empty(&queue)) {
        Vec2i point = queue_pop(&queue);
        i32 dist = map.dist[point.y * map.width + point.x];
        for (i32 i = 0; i < 4; i++) {
            i32 nx = point.x + DX[i];
            i32 ny = point.y + DY[i];
            if (nx >= 0 && ny >= 0 && nx < map.width && ny < map.height) {
                u8 cell = grid.grid[ny * map.width + nx];
                if (cell != '#') {
                    if (dist + 1 < map.dist[ny * map.width + nx]) {
                        map.dist[ny * map.width + nx] = dist + 1;
                        queue_push(&queue, vec2i(nx, ny));
                    }
                }
            }
        }
    }
    return map;
}

static void
dump_distance_map(DistanceMap map) {
    for (i32 y = 0; y < map.height; y++) {
        for (i32 x = 0; x < map.width; x++) {
            i32 dist = map.dist[y * map.width + x];
            if (dist == INT32_MAX) {
                printf("#");
            } else {
                printf("%d", dist % 10);
            }
        }
        printf("\n");
    }
}

typedef DYNAMIC_ARRAY(Cheat) Cheats;

static Cheats
enumerate_cheats(Arena *arena, Grid *grid, i32 length) {
    Cheats cheats = {0};
    i32 radius = length; // manhattan distance between start and end
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            /* printf("@ (%d, %d)\n", x, y); */
            for (i32 dy = -radius; dy <= radius; dy++) {
                // abs(dx) + abs(dy) <= radius
                for (i32 dx = -radius; dx <= radius; dx++) {
                    i32 d = abs(dx) + abs(dy);
                    if (d == 0 || d > radius) {
                        continue;
                    }

                    i32 x_ = x + dx;
                    i32 y_ = y + dy;

                    if (x_ < 0 || y_ < 0 || x_ >= grid->width || y_ >= grid->height) {
                        continue;
                    }

                    i32 source_index = y * grid->width + x;
                    i32 target_index = y_ * grid->width + x_;
                    if (grid->grid[source_index] != '#' && grid->grid[target_index] != '#') {
                        *push(&cheats, arena) = (Cheat) { .start_x = x, .start_y = y, .end_x = x_, .end_y = y_ };
                        /* i32 len = abs(dx) + abs(dy); */
                        /* printf("Found cheat of length %d between (%d, %d) and (%d, %d)\n", len - 1, x, y, x_, y_); */
                    }
                }
            }
        }
    }
    return cheats;
}

static i32
count_savings(Arena *arena, Grid grid, i32 start_x, i32 start_y, i32 end_x, i32 end_y, DistanceMap dm_start, DistanceMap dm_end, i32 cheat_duration) {
    Cheats cheats = enumerate_cheats(arena, &grid, cheat_duration);
    /* printf("Found %ld possible cheats\n", cheats.len); */

    i32 without_cheats = dm_end.dist[start_y * grid.width + start_x];
    /* printf("without cheats: %d\n", without_cheats); */

    i32 count = 0;
    for (isize i = 0; i < cheats.len; i++) {
        Cheat cheat = cheats.data[i];
        ASSERT(grid.grid[cheat.start_y * grid.width + cheat.start_x] != '#');
        ASSERT(grid.grid[cheat.end_y * grid.width + cheat.end_x] != '#');
        i32 dx = cheat.end_x - cheat.start_x;
        i32 dy = cheat.end_y - cheat.start_y;
        i32 cheat_len = abs(dx) + abs(dy);
        ASSERT(cheat_len <= cheat_duration);
        ASSERT(cheat_len > 0);

        i32 start_index = cheat.start_y * grid.width + cheat.start_x;
        i32 end_index   = cheat.end_y   * grid.width + cheat.end_x;

        i32 saved = without_cheats - (dm_start.dist[start_index] + cheat_len + dm_end.dist[end_index]);
        bool bar = false;
        if (saved > 0) {
#if 0
            printf("From (%d, %d) [%d] via cheat (%d, %d) -> (%d, %d) [%d] to (%d, %d) [%d]\n",
                    start_x, start_y, dm_start.dist[start_index],
                    cheat.start_x, cheat.start_y, cheat.end_x, cheat.end_y, cheat_len,
                    end_x, end_y, dm_end.dist[end_index]);
#endif
#if 0
            printf("Saved %d\n", saved);
#endif
        }
        if (saved >= 100) count++;
    }
    return count;
}

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

    i32 start_x = -1, start_y = -1;
    i32 end_x = -1, end_y = -1;

    for (i32 y = 0; y < grid.height; y++) {
        for (i32 x = 0; x < grid.width; x++) {
            if (grid_at(&grid, x, y) == 'S') {
                start_x = x;
                start_y = y;
            }
            else if (grid_at(&grid, x, y) == 'E') {
                end_x = x;
                end_y = y;
            }
        }
    }
    ASSERT(start_x != -1 && start_y != -1);
    ASSERT(end_x != -1 && end_y != -1);

    grid.grid[start_y * grid.width + start_x] = '.';

    DistanceMap dm_start = compute_distances(arena, grid, start_x, start_y);
    /* dump_distance_map(dm_start); */
    DistanceMap dm_end = compute_distances(arena, grid, end_x, end_y);
    /* dump_distance_map(dm_end); */
    ASSERT(dm_start.dist[1 * grid.width + 3] + dm_end.dist[1 * grid.width + 3] == 
            dm_end.dist[start_y * grid.width + start_x]);

    Arena watermark = *arena;
    i32 part_1 = count_savings(arena, grid, start_x, start_y, end_x, end_y, dm_start, dm_end, 2);
    printf("%d\n", part_1);

    *arena = watermark;
    i32 part_2 = count_savings(arena, grid, start_x, start_y, end_x, end_y, dm_start, dm_end, 20);
    printf("%d\n", part_2);
}