const std = @import("std");

const ArrayList = std.ArrayList;
const PathQueue = std.PriorityQueue(Path, void, sort_queue);
const TestMap = std.AutoHashMap(SteppedOn, bool);

const SteppedOn = struct {
    x: i32 = 0,
    y: i32 = 0,
    d_x: i32 = 0,
    d_y: i32 = 0,
};

const Path = struct {
    x: i32 = 0,
    y: i32 = 0,
    d_x: i32 = 0,
    d_y: i32 = 0,
    heat_loss: u32 = 0,
    steps_left: i32 = 0,
    history: ArrayList([5]i32),
};

pub fn solve_part1(allocator: std.mem.Allocator, file: []const u8) !void {
    var inputs = try std.fs.cwd().openFile(file, .{});

    var contents = try inputs.readToEndAlloc(allocator, 100 * 1024);

    var grid_map = ArrayList([]u8).init(allocator);
    var walked = TestMap.init(allocator);

    const print_map_1 = false;
    const print_map_2 = true;

    var last_index: u64 = 0;

    while (std.mem.indexOfScalar(u8, contents[last_index..], '\n')) |new_line_index| {
        try grid_map.append(contents[last_index .. last_index + new_line_index]);
        last_index += new_line_index + 1;

        if (last_index >= contents.len) break;
    }

    var queue = PathQueue.init(allocator, {});

    try queue.add(.{ .x = 0, .y = 0, .d_x = 1, .d_y = 0, .heat_loss = 0, .steps_left = 2, .history = ArrayList([5]i32).init(allocator) });
    try queue.add(.{ .x = 0, .y = 0, .d_x = 0, .d_y = 1, .heat_loss = 0, .steps_left = 2, .history = ArrayList([5]i32).init(allocator) });

    var final_heat_loss: u32 = 0;

    while (queue.removeOrNull()) |path| {
        if (try walk(path, &queue, grid_map, &walked, print_map_1, false)) |d| {
            if (d.heat_loss >= 0) {
                final_heat_loss = @intCast(d.heat_loss);
                if (print_map_1) {
                    var tt = try allocator.alloc([]u8, grid_map.items.len);

                    for (0..tt.len) |l_i| {
                        tt[l_i] = try allocator.alloc(u8, grid_map.items[0].len);
                        @memset(tt[l_i], '0');
                    }

                    for (d.history.items) |hh| {
                        tt[@intCast(hh[1])][@intCast(hh[0])] = grid_map.items[@intCast(hh[1])][@intCast(hh[0])];
                    }

                    for (tt) |line| {
                        std.log.info("{s}", .{line});
                    }
                }
                break;
            }
        }
    }
    std.log.info("Part 1: {d}", .{final_heat_loss});

    queue.deinit();
    queue = PathQueue.init(allocator, {});
    try queue.add(.{ .x = 0, .y = 0, .d_x = 1, .d_y = 0, .heat_loss = 0, .steps_left = 9, .history = ArrayList([5]i32).init(allocator) });
    try queue.add(.{ .x = 0, .y = 0, .d_x = 0, .d_y = 1, .heat_loss = 0, .steps_left = 9, .history = ArrayList([5]i32).init(allocator) });

    final_heat_loss = 0;

    walked.clearRetainingCapacity();

    while (queue.removeOrNull()) |path| {
        if (try walk(path, &queue, grid_map, &walked, print_map_2, true)) |d| {
            if (d.heat_loss >= 0) {
                final_heat_loss = @intCast(d.heat_loss);

                if (print_map_2) {
                    var tt = try allocator.alloc([]u8, grid_map.items.len);
                    defer {
                        for (tt) |line| {
                            allocator.free(line);
                        }

                        allocator.free(tt);
                    }

                    for (0..tt.len) |l_i| {
                        tt[l_i] = try allocator.alloc(u8, grid_map.items[0].len);
                        @memset(tt[l_i], '0');
                    }

                    for (d.history.items) |hh| {
                        tt[@intCast(hh[1])][@intCast(hh[0])] = grid_map.items[@intCast(hh[1])][@intCast(hh[0])];
                    }

                    for (tt) |line| {
                        std.log.info("{s}", .{line});
                    }
                }
                break;
            }
        }
    }

    std.log.info("Part 2: {d}", .{final_heat_loss});
}

fn sort_queue(context: void, a: Path, b: Path) std.math.Order {
    _ = context;

    return std.math.order(a.heat_loss, b.heat_loss);
}

fn walk(path: Path, queue: *PathQueue, grid: ArrayList([]u8), walked: *TestMap, history: bool, part2: bool) !?Path {
    var width = grid.items[0].len;
    const required_steps_left = 6;
    var height = grid.items.len;

    if (path.x + path.d_x >= width or path.x + path.d_x < 0) {
        return null;
    }
    if (path.y + path.d_y >= height or path.y + path.d_y < 0) {
        return null;
    }

    var new_pos = path;
    new_pos.x += new_pos.d_x;
    new_pos.y += new_pos.d_y;
    if (history) try new_pos.history.append([5]i32{ new_pos.x, new_pos.y, new_pos.d_x, new_pos.d_y, @intCast(new_pos.heat_loss) });
    var p_x: u32 = @intCast(new_pos.x);
    var p_y: u32 = @intCast(new_pos.y);
    new_pos.heat_loss += grid.items[p_y][p_x] - '0';

    var map_key = SteppedOn{ .x = new_pos.x, .y = new_pos.y, .d_x = new_pos.d_x, .d_y = new_pos.d_y };

    if (new_pos.steps_left > 0) {
        new_pos.steps_left -= 1;
        try queue.add(new_pos);
    }
    if (new_pos.x == width - 1 and new_pos.y == height - 1 and !(part2 and path.steps_left > required_steps_left)) return new_pos;

    if (!walked.contains(map_key)) {
        if (new_pos.steps_left < required_steps_left) {
            try walked.put(map_key, true);
            new_pos.steps_left = if (part2) 9 else 2;
            new_pos.d_x = path.d_y * 1;
            new_pos.d_y = path.d_x * 1;
            if (history) new_pos.history = try new_pos.history.clone();
            try queue.add(new_pos);

            new_pos.d_x *= -1;
            new_pos.d_y *= -1;
            if (history) new_pos.history = try new_pos.history.clone();
            try queue.add(new_pos);
        }
    }

    return null;
}