const std = @import("std");

pub fn main() !void {
    const content = @embedFile("data.txt");

    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    var grid = std.ArrayList(std.ArrayList(u8)).init(allocator);
    defer grid.deinit();

    var grid2 = std.ArrayList(std.ArrayList(u8)).init(allocator);
    defer grid2.deinit();

    var lines = std.mem.split(u8, content, "\n");

    while (lines.next()) |line| {
        if (line.len < 1) {
            break;
        }

        var lineArrayList = std.ArrayList(u8).init(allocator);
        var lineArrayList2 = std.ArrayList(u8).init(allocator);

        for (line) |char| {
            try lineArrayList.append(char);
            try lineArrayList2.append(char);
        }

        try grid.append(lineArrayList);
        try grid2.append(lineArrayList2);
    }

    var map = std.StringHashMap(usize).init(allocator);
    defer map.deinit();

    transpose(&grid);
    try tiltWest(&grid);
    transpose(&grid);

    var part1 = calculateLoad(grid);

    std.debug.print("part 1: {d}\n", .{part1});

    var prev: usize = 0;
    var loop: usize = 0;

    const MAX = 1_000_000_000;
    for (0..MAX) |i| {
        // see if this grid arrangement is in our map
        const key = try getKey(&grid2);
        const hashed = map.get(key);

        // if it is, we've found a loop
        if (hashed) |val| {
            prev = val;
            loop = i;
            break;
        }

        try cycle(&grid2);

        // add the new grid arrangement to the map
        try map.put(key, i);
    }

    // figure out where in the loop we should stop
    const stop = (MAX - prev) % (loop - prev);

    // cycle that many more times
    for (0..stop) |_| {
        try cycle(&grid2);
    }
    var part2 = calculateLoad(grid2);

    std.debug.print("part 2: {d}\n", .{part2});
}

fn getKey(grid: *std.ArrayList(std.ArrayList(u8))) ![]u8 {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    var str = std.ArrayList(u8).init(allocator);
    for (grid.items) |row| {
        try str.appendSlice(row.items);
    }
    return str.items;
}

fn cycle(grid: *std.ArrayList(std.ArrayList(u8))) !void {
    transpose(grid); // to tilt north, transpose so north is left
    try tiltWest(grid);

    transpose(grid); // to tilt west, transpose again

    try tiltWest(grid);
    // at this point, grid is back to normal
    flipY(grid); // to tilt south, flip (south is up) then transpose (south is left)
    transpose(grid);

    try tiltWest(grid);
    // at this point, south is left, west is up. we want east left
    transpose(grid); // south is up, west is left
    flipY(grid); // grid is back to normal
    flipX(grid);
    try tiltWest(grid);
    flipX(grid);
}

fn calculateLoad(grid: std.ArrayList(std.ArrayList(u8))) usize {
    var load: usize = 0;
    for (grid.items, 0..) |row, rowNum| {
        var rowMultiplier = grid.items.len - rowNum;
        for (row.items) |cell| {
            if (cell == 'O') {
                load += rowMultiplier;
            }
        }
    }

    return load;
}

fn tiltWest(grid: *std.ArrayList(std.ArrayList(u8))) !void {
    for (grid.items) |row| {
        for (row.items, 0..) |cell, i| {
            if (cell == 'O' and i != 0 and row.items[i - 1] == '.') {
                var j = i;
                while (j > 0) {
                    row.items[j - 1] = 'O';
                    row.items[j] = '.';
                    j -= 1;
                    if (j < 1 or row.items[j - 1] != '.') {
                        break;
                    }
                }
            }
        }
    }
}

// assumes square grid
fn transpose(grid: *std.ArrayList(std.ArrayList(u8))) void {
    const length = grid.items.len;

    for (0..length - 1) |i| {
        for (i + 1..length) |j| {
            const temp = grid.items[i].items[j];
            grid.items[i].items[j] = grid.items[j].items[i];
            grid.items[j].items[i] = temp;
        }
    }
}

// flip a grid upside down
fn flipY(grid: *std.ArrayList(std.ArrayList(u8))) void {
    const height = grid.items.len;

    for (0..height / 2) |r| {
        const temp = grid.items[r];
        grid.items[r] = grid.items[height - r - 1];
        grid.items[height - r - 1] = temp;
    }
}

// flip a grid left to right
fn flipX(grid: *std.ArrayList(std.ArrayList(u8))) void {
    const length = grid.items.len;

    for (0..length) |r| {
        for (0..length / 2) |c| {
            const temp = grid.items[r].items[c];
            grid.items[r].items[c] = grid.items[r].items[length - c - 1];
            grid.items[r].items[length - c - 1] = temp;
        }
    }
}