const std = @import("std");

const Pos = @Vector(2, usize);

const Dir = enum {
    right,
    up,
    left,
    down,

    fn rotate(self: @This()) @This() {
        return switch (self) {
            .right => .down,
            .up => .right,
            .left => .up,
            .down => .left,
        };
    }
};

const Map = struct {
    obstacles: std.AutoHashMap(Pos, void),
    size: Pos,

    const Self = @This();

    fn step_forwards(self: Self, pos: Pos, dir: Dir) ?Pos {
        return switch (dir) {
            .right => .{ if (pos[0] + 1 == self.size[0]) return null else pos[0] + 1, pos[1] },
            .up => .{ pos[0], std.math.sub(usize, pos[1], 1) catch return null },
            .left => .{ std.math.sub(usize, pos[0], 1) catch return null, pos[1] },
            .down => .{ pos[0], if (pos[1] + 1 == self.size[1]) return null else pos[1] + 1 },
        };
    }

    fn step(self: Self, pos: Pos, dir: Dir) ?struct { Pos, Dir } {
        var new_pos = self.step_forwards(pos, dir) orelse return null;
        var new_dir = dir;
        while (self.obstacles.contains(new_pos)) {
            new_dir = new_dir.rotate();
            new_pos = self.step_forwards(pos, new_dir) orelse return null;
        }
        return .{ new_pos, new_dir };
    }
};

const Input = struct {
    map: Map,
    start_pos: Pos,
    start_dir: Dir,
};

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var lines = std.mem.split(u8, data, "\n");
    var input = Input{
        .map = .{
            .obstacles = std.AutoHashMap(Pos, void).init(allocator),
            .size = undefined,
        },
        .start_pos = undefined,
        .start_dir = undefined,
    };
    var y: usize = 0;
    while (lines.next()) |line| : (y += 1) {
        if (line.len == 0) input.map.size[1] = y else input.map.size[0] = line.len;
        for (0.., line) |x, c| {
            if (c == '#') {
                try input.map.obstacles.put(.{ x, y }, {});
            } else if (c == '>') {
                input.start_pos = .{ x, y };
                input.start_dir = .right;
            } else if (c == '^') {
                input.start_pos = .{ x, y };
                input.start_dir = .up;
            } else if (c == '<') {
                input.start_pos = .{ x, y };
                input.start_dir = .left;
            } else if (c == 'v') {
                input.start_pos = .{ x, y };
                input.start_dir = .down;
            }
        }
    }
    return input;
}

fn printAndWait(map: Map, traversed: std.AutoHashMap(Pos, void)) void {
    var underlying_buffer: [100000]u8 = undefined;
    var buffer = std.io.fixedBufferStream(&underlying_buffer);
    var writer = buffer.writer();
    for (0..map.size[1]) |y| {
        for (0..map.size[0]) |x| {
            writer.print("{c}", .{
                if (traversed.contains(.{ x, y })) 'X' else if (map.obstacles.contains(.{ x, y })) @as(u8, '#') else '.',
            }) catch unreachable;
        }
        writer.print("\n", .{}) catch unreachable;
    }
    var keys = traversed.keyIterator();
    while (keys.next()) |pos| {
        if (@reduce(.Or, pos.* >= map.size)) {
            writer.print("oob: {}\n", .{pos.*}) catch unreachable;
        }
    }
    std.debug.print("{s}\n", .{underlying_buffer[0..try buffer.getPos()]});
    // try std.io.getStdIn().reader().skipUntilDelimiterOrEof('\n');
}

const test_input =
    \\....#.....
    \\.........#
    \\..........
    \\..#.......
    \\.......#..
    \\..........
    \\.#..^.....
    \\........#.
    \\#.........
    \\......#...
    \\
;

test "part1" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part1(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 41);
}

pub fn part1(allocator: std.mem.Allocator, input: Input) !u32 {
    var pos = input.start_pos;
    var dir = input.start_dir;
    var traversed = std.AutoHashMap(Pos, void).init(allocator);
    defer printAndWait(input.map, traversed);
    return while (true) {
        try traversed.put(pos, {});
        const new = input.map.step(pos, dir);
        if (new) |n| {
            pos = n[0];
            dir = n[1];
            // try printAndWait(input.map, traversed);
        } else {
            break traversed.count();
        }
    };
}

test "part2" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();

    const output = try part2(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 6);
}

pub fn part2(allocator: std.mem.Allocator, input: Input) !u32 {
    var pos = input.start_pos;
    var dir = input.start_dir;
    var options = std.AutoHashMap(Pos, void).init(allocator);
    var traversed = std.AutoHashMap(Pos, void).init(allocator);
    var map: Map = .{ .obstacles = try input.map.obstacles.clone(), .size = input.map.size };
    return while (true) {
        try traversed.put(pos, {});
        const next = map.step(pos, dir) orelse break options.count();
        defer pos = next[0];
        defer dir = next[1];
        const next_pos = next[0];
        if (map.obstacles.contains(next_pos) or traversed.contains(next_pos)) continue;
        try map.obstacles.put(next_pos, {});
        if (try gets_stuck(allocator, map, pos, dir)) {
            try options.put(next_pos, {});
        }
        std.debug.assert(map.obstacles.remove(next_pos));
    };
}

fn gets_stuck(allocator: std.mem.Allocator, map: Map, start_pos: Pos, start_dir: Dir) !bool {
    var pos = start_pos;
    var dir = start_dir;
    var traversed = std.AutoHashMap(struct { Pos, Dir }, void).init(allocator);
    defer traversed.deinit();
    return while (true) {
        const res = try traversed.getOrPut(.{ pos, dir });
        if (res.found_existing) break true;
        const new = map.step(pos, dir);
        if (new) |n| {
            pos = n[0];
            dir = n[1];
            // try printAndWait(input.map, traversed);
        } else {
            break false;
        }
    };
}