#include "aoc.h"
#include "limits.h"

typedef struct {
    i32 x, y;
} Vec2i;

typedef enum {
    Right,
    Down,
    Left,
    Up,
} Orientation;

static inline Orientation
turn(Orientation orientation, i32 delta) {
    return (orientation + delta) % 4;
}

static const Vec2i OFFSET[] = {
    [Right] = { .x = 1, .y = 0 },
    [Down]  = { .x = 0, .y = 1 },
    [Left]  = { .x = -1, .y = 0 },
    [Up]    = { .x = 0, .y = -1 },
};

typedef struct {
    i32 x, y;
    Orientation orientation;
} Tile;

typedef struct TileQueue {
    Tile *tiles;
    isize len;
    isize read, write;
} TileQueue;

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

static Tile
tile_queue_pop(TileQueue *queue) {
    ASSERT(!tile_queue_empty(queue));
    Tile result = queue->tiles[queue->read];
    queue->read = (queue->read + 1) % queue->len;
    return result;
}

static void
tile_queue_push(TileQueue *queue, i32 x, i32 y, Orientation orientation) {
    isize next = (queue->write + 1) % queue->len;
    ASSERT(next != queue->read);
    queue->tiles[queue->write] = (Tile) { .x = x, .y = y, .orientation = orientation };
    queue->write = next;
}

typedef struct {
    i32 distance;
    DYNAMIC_ARRAY(Tile) predecessors;
} Cell;

static void
solve(Arena *arena, Grid *grid) {
    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->grid[y * grid->width + x] == 'E') {
                end_x = x;
                end_y = y;
            }
            else if (grid->grid[y * grid->width + x] == 'S') {
                start_x = x;
                start_y = y;
            }
        }
    }
    ASSERT(start_x != -1 && start_y != -1);
    ASSERT(end_x != -1 && end_y != -1);

    Cell *cells = ARENA_ALLOC_ARRAY(arena, Cell, 4 * grid->width * grid->height);
    for (isize i = 0; i < 4 * grid->width * grid->height; i++) {
        cells[i].distance = INT_MAX;
    }
    // `cells` is a 3d array indexed as cells[4 * (y * width + x) + orientation]
#define CELL(x, y, o) cells[4 * ((y) * grid->width + (x)) + (o)]
#define DISTANCE(x, y, o) CELL(x, y, o).distance

    Orientation start_orientation = Right;
    DISTANCE(start_x, start_y, start_orientation) = 0;
    isize queue_size = 2048;
    TileQueue queue = {
        .tiles = ARENA_ALLOC_ARRAY(arena, Tile, queue_size),
        .len = queue_size,
    };
    tile_queue_push(&queue, start_x, start_y, start_orientation);

    while (!tile_queue_empty(&queue)) {
        // walk straight
        Tile current = tile_queue_pop(&queue);
        i32 x = current.x;
        i32 y = current.y;
        Orientation orientation = current.orientation;
        i32 dist = DISTANCE(current.x, current.y, current.orientation);
        if (grid->grid[y * grid->width + x] == '#') {
            continue;
        }
        i32 nx = current.x + OFFSET[current.orientation].x;
        i32 ny = current.y + OFFSET[current.orientation].y;
#define TRY(x_, y_, o_, delta) do { \
    i32 candidate = DISTANCE(x_, y_, o_); \
    if (dist + delta <= candidate) { \
        DISTANCE(x_, y_, o_) = dist + delta; \
        Cell *cell = &CELL(x_, y_, o_); \
        if (dist + delta < candidate) { \
            tile_queue_push(&queue, x_, y_, o_); \
            cell->predecessors.len = 0; \
        } \
        *push(&cell->predecessors, arena) = (Tile) { .x = x, .y = y, .orientation = orientation }; \
    } \
} while (0)
        TRY(nx, ny, orientation, 1);
        // turn 90deg cw
        TRY(x, y, turn(orientation, 1), 1000);
        // turn 180deg cw
        TRY(x, y, turn(orientation, 2), 2000);
        // turn 90deg ccw
        TRY(x, y, turn(orientation, 3), 1000);
#undef TRY
    }

    i32 part_1 = MIN4(DISTANCE(end_x, end_y, Right), DISTANCE(end_x, end_y, Down), DISTANCE(end_x, end_y, Left), DISTANCE(end_x, end_y, Up));
    
    // Reconstruct best path
    u8 *best_path = ARENA_ALLOC_ARRAY(arena, u8, grid->width * grid->height);
    queue.write = 0;
    queue.read = 0;
    tile_queue_push(&queue, end_x, end_y, Right);
    tile_queue_push(&queue, end_x, end_y, Down);
    tile_queue_push(&queue, end_x, end_y, Left);
    tile_queue_push(&queue, end_x, end_y, Up);
    while (!tile_queue_empty(&queue)) {
        Tile tile = tile_queue_pop(&queue);
        best_path[tile.y * grid->width + tile.x] = 1;
        Cell *cell = &CELL(tile.x, tile.y, tile.orientation);
        for (isize i = 0; i < cell->predecessors.len; i++) {
            Tile pred = cell->predecessors.data[i];
            tile_queue_push(&queue, pred.x, pred.y, pred.orientation);
        }
    }

    i32 part_2 = 0;
    for (i32 y = 0; y < grid->height; y++) {
        for (i32 x = 0; x < grid->width; x++) {
            part_2 += (i32) best_path[y * grid->width + x];
        }
    }

    printf("%d\n", part_1);
    printf("%d\n", part_2);
#undef DISTANCE
#undef CELL
}

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

    solve(arena, &input);
}