const std = @import("std");
const mecha = @import("mecha");
const print = std.debug.print;

pub fn solve(part: []u8, buffer: []u8, allocator: std.mem.Allocator) !void {
    var lines = std.mem.splitScalar(u8, buffer, '\n');
    var grid = std.ArrayList([]u8).init(allocator);
    defer grid.deinit();
    var x: i32 = 0;
    var y: i32 = 0;

    var row: usize = 0;
    while (lines.next()) |line| : (row += 1) {
        for (line, 0..) |element, i| {
            if (element == '^') {
                x = @intCast(i);
                y = @intCast(row);
            }
        }
        try grid.append(@constCast(line));
    }

    if (std.mem.eql(u8, part, "1")) {
        try part1(&grid, x, y, allocator);
    } else {
        try part2(&grid, x, y, allocator);
    }
}

fn locations(grid: *std.ArrayList([]u8), startx: i32, starty: i32, allocator: std.mem.Allocator) !std.AutoHashMap(std.meta.Tuple(&.{ i32, i32, i32, i32 }), void) {
    var dirx: i32 = 0;
    var diry: i32 = -1;
    var x = startx;
    var y = starty;

    var history = try allocator.alloc(bool, grid.items.len * grid.items[0].len * 4);
    defer allocator.free(history);

    for (history) |*elem| {
        elem.* = false;
    }

    var unique = std.AutoHashMap(std.meta.Tuple(&.{ i32, i32, i32, i32 }), void).init(allocator);
    errdefer unique.deinit();
    const width: i32 = @intCast(grid.items[0].len);

    while (y >= 0 and y < grid.items.len and x >= 0 and x < grid.items[@intCast(y)].len) {
        if (!history[@intCast(y * width + x)] and grid.items[@intCast(y)][@intCast(x)] != '#') {
            try unique.put(.{ x, y, dirx, diry }, {});
            history[@intCast(y * width + x)] = true;
        }

        if (grid.items[@intCast(y)][@intCast(x)] == '#') {
            x -= dirx;
            y -= diry;
            const temp = dirx;
            dirx = -diry;
            diry = temp;
        }

        x += dirx;
        y += diry;
    }

    return unique;
}

fn part1(grid: *std.ArrayList([]u8), startx: i32, starty: i32, allocator: std.mem.Allocator) !void {
    var locs = try locations(grid, startx, starty, allocator);
    defer locs.deinit();
    print("{d}\n", .{locs.count()});
}

//----------------------------PART 2---------------------------------------------

fn directionToIndex(dirx: i32, diry: i32) i32 {
    if (dirx == 0 and diry == -1) {
        return 0;
    } else if (dirx == 1 and diry == 0) {
        return 1;
    } else if (dirx == 0 and diry == 1) {
        return 2;
    } else {
        return 3;
    }
}
fn part2(grid: *std.ArrayList([]u8), startx: i32, starty: i32, allocator: std.mem.Allocator) !void {
    var locs = try locations(grid, startx, starty, allocator);
    defer locs.deinit();
    var total: usize = 0;

    var iterator = locs.keyIterator();
    var history = try allocator.alloc(bool, grid.items.len * grid.items[0].len * 4);
    defer allocator.free(history);

    const width: i32 = @intCast(grid.items[0].len);

    while (iterator.next()) |location| {
        grid.items[@intCast(location[1])][@intCast(location[0])] = '#';
        var dirx: i32 = location[2];
        var diry: i32 = location[3];
        var x = location[0];
        var y = location[1];

        for (history) |*elem| {
            elem.* = false;
        }

        while (y >= 0 and y < grid.items.len and x >= 0 and x < grid.items[@intCast(y)].len) {
            const index = y * width * 4 + x * 4 + directionToIndex(dirx, diry);
            if (history[@intCast(index)]) {
                total += 1;
                break;
            }
            if (grid.items[@intCast(y)][@intCast(x)] == '#') {
                x -= dirx;
                y -= diry;
                const temp = dirx;
                dirx = -diry;
                diry = temp;
            }

            history[@intCast(index)] = true;

            x += dirx;
            y += diry;
        }

        grid.items[@intCast(location[1])][@intCast(location[0])] = '.';
    }

    print("{d}\n", .{total});
}