const std = @import("std");
const aoc = @import("./aoc.zig");
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;

const Tile = enum {
    Empty,
    Rock,
    Wall
};

const Input = struct {
    tiles: std.ArrayList(Tile),
    width: u32,
    height: u32,

    fn init(allocator: Allocator) Input {
        return Input{
            .tiles = std.ArrayList(Tile).init(allocator),
            .width = 0,
            .height = 0,
        };
    }

    fn deinit(self: Input) void {
        self.tiles.deinit();
    }

    fn getIndex(self: Input, x: u32, y: u32) usize {
        assert(x < self.width);
        assert(y < self.height);
        return y * self.width + x;
    }

    fn get(self: Input, x: u32, y: u32) Tile {
        return self.tiles.items[self.getIndex(x, y)];
    }

    fn set(self: Input, x: u32, y: u32, tile: Tile) void {
        self.tiles.items[self.getIndex(x, y)] = tile;
    }

    fn inBounds(self: Input, x: i32, y: i32) bool {
        return (0 <= x and x < self.width) and (0 <= y and y < self.height);
    }

    fn print(self: Input) void {
        for (0..self.height) |y| {
            for (0..self.height) |x| {
                const idx = y * self.width + x;
                const tile = self.tiles.items[idx];
                const symbol: u8 = switch (tile) {
                    .Empty => '.',
                    .Wall => '#',
                    .Rock => 'O'
                };

                std.debug.print("{c}", .{symbol});
            }
            std.debug.print("\n", .{});
        }
    }
};

fn parseInput(allocator: Allocator, lines: []const []const u8) !Input {
    var parsed = Input.init(allocator);
    errdefer parsed.deinit();
    parsed.width = @intCast(lines[0].len);
    parsed.height = @intCast(lines.len);

    try parsed.tiles.ensureTotalCapacity(parsed.width * parsed.height);

    for (lines) |line| {
        for (line) |symbol| {
            const tile = switch (symbol) {
                'O' => Tile.Rock,
                '.' => Tile.Empty,
                '#' => Tile.Wall,
                else => return error.InvalidTile
            };

            try parsed.tiles.append(tile);
        }
    }

    return parsed;
}

fn slideRock(input: Input, x: u32, y: u32, dir_x: i32, dir_y: i32) void {
    assert(input.get(x, y) == .Rock);

    var new_x = x;
    var new_y = y;
    while (true) {
        const next_x = @as(i32, @intCast(new_x)) + dir_x;
        const next_y = @as(i32, @intCast(new_y)) + dir_y;
        if (!input.inBounds(next_x, next_y)) break;
        if (input.get(@intCast(next_x), @intCast(next_y)) != .Empty) break;

        new_x = @intCast(next_x);
        new_y = @intCast(next_y);
    }

    input.set(x, y, .Empty);
    input.set(new_x, new_y, .Rock);
}

fn slideRocks(input: Input, dir_x: i32, dir_y: i32) void {
    if (dir_y != 0) {
        // Column-based
        for (0..input.height) |oy| {
            for (0..input.width) |ox| {
                const x: u32 = @intCast(ox);
                const y: u32 = @intCast(@as(i32, @intCast(oy)) * (-dir_y) + @divExact(dir_y + 1, 2) * @as(i32, @intCast(input.height-1)));
                if (input.get(x, y) == .Rock) {
                    slideRock(input, x, y, dir_x, dir_y);
                }
            }
        }
    } else {
        // Row-based
        for (0..input.width) |ox| {
            for (0..input.height) |oy| {
                const x: u32 = @intCast(@as(i32, @intCast(ox)) * (-dir_x) + @divExact(dir_x + 1, 2) * @as(i32, @intCast(input.width-1)));
                const y: u32 = @intCast(oy);
                if (input.get(x, y) == .Rock) {
                    slideRock(input, x, y, dir_x, dir_y);
                }
            }
        }
    }
}

fn cycleSlide(input: Input) void {
    const dirs = &.{
        .{ .x =  0, .y = -1 }, // North
        .{ .x = -1, .y =  0 }, // West
        .{ .x =  0, .y =  1 }, // South
        .{ .x =  1, .y =  0 }, // East
    };

    inline for (dirs) |dir| {
        slideRocks(input, dir.x, dir.y);
    }
}

fn getTotalLoad(input: Input) u64 {
    var result: u64 = 0;
    for (0..input.height) |y| {
        var count: u32 = 0;
        for (0..input.width) |x| {
            if (input.get(@intCast(x), @intCast(y)) == .Rock) {
                count += 1;
            }
        }
        result += count * (input.height - y);
    }

    return result;
}

pub fn part1(input: *aoc.Input) !aoc.Result {
    const parsed = try parseInput(input.allocator, input.lines);
    defer parsed.deinit();

    slideRocks(parsed, 0, -1);

    return .{ .uint = getTotalLoad(parsed) };
}

pub fn part2(input: *aoc.Input) !aoc.Result {
    const target_cycle_count = 1000000000;

    const allocator = input.allocator;
    const parsed = try parseInput(allocator, input.lines);
    defer parsed.deinit();

    var seen = std.ArrayList(std.ArrayList(Tile)).init(allocator);
    defer {
        for (seen.items) |tiles| {
            tiles.deinit();
        }
        seen.deinit();
    }

    var cycle_size: u32 = 0;
    var cycle_start: u32 = 0;
    var iteration: u32 = 0;
    outer: while (true) {
        cycleSlide(parsed);

        for (0.., seen.items) |i, tiles| {
            if (std.mem.eql(Tile, tiles.items, parsed.tiles.items)) {
                cycle_start = @intCast(i);
                cycle_size = @intCast(iteration - i);
                break :outer;
            }
        }
        try seen.append(try parsed.tiles.clone());

        iteration += 1;
    }

    const skipped_map = Input{
        .tiles = seen.items[@mod((target_cycle_count-1) - cycle_start, cycle_size) + cycle_start],
        .width = parsed.width,
        .height = parsed.height
    };

    return .{ .uint = getTotalLoad(skipped_map) };
}

const example_input = [_][]const u8{
    "O....#....",
    "O.OO#....#",
    ".....##...",
    "OO.#O....O",
    ".O.....O#.",
    "O.#..O.#.#",
    "..O..#O..O",
    ".......O..",
    "#....###..",
    "#OO..#....",
};

test "part 1 example" {
    try aoc.expectAnswerUInt(part1, 136, &example_input);
}

test "part 2 example" {
    try aoc.expectAnswerUInt(part2, 64, &example_input);
}