const std = @import("std");

const Galaxy = struct { row: usize, col: usize };

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

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

    // put into ArrayList of ArrayLists
    var grid = try parseGrid(content, allocator);
    defer grid.deinit();
    std.debug.print("grid: {any}\n", .{grid});

    // find empty rows
    const emptyRows = try findEmptyRows(grid, allocator);
    std.debug.print("empty rows: {any}\n", .{emptyRows});

    // find empty columns
    const emptyColumns = try findEmptyColumns(grid, allocator);
    std.debug.print("empty cols: {any}\n", .{emptyColumns});

    // duplicate empty rows
    // NOTE: not doing this any more, using Part 2 solution for Part 1
    //
    // for (emptyRows.items, 0..) |row, i| {
    //     try grid.insert(row + i, grid.items[row + i]);
    // }
    //
    // // duplicate empty columns
    // for (emptyColumns.items, 0..) |col, i| {
    //     for (grid.items, 0..) |_, row| {
    //         try grid.items[row].insert(col + i, '.');
    //     }
    // }

    // find original galaxy locations
    var galaxies = std.ArrayList(Galaxy).init(allocator);
    for (grid.items, 0..) |row, i| {
        for (row.items, 0..) |char, j| {
            if (char == '#') {
                try galaxies.append(Galaxy{ .row = i, .col = j });
            }
        }
    }

    var galaxies1 = std.ArrayList(Galaxy).init(allocator);
    var galaxies2 = std.ArrayList(Galaxy).init(allocator);

    for (galaxies.items) |g| {
        const row = g.row;
        const col = g.col;

        // how many expanded rows & cols before this galaxy?
        var expandedRows: usize = 0;
        for (emptyRows.items) |r| {
            if (r > row) {
                break;
            }
            expandedRows += 1;
        }

        var expandedColumns: usize = 0;
        for (emptyColumns.items) |c| {
            if (c > col) {
                break;
            }
            expandedColumns += 1;
        }

        const expansionFactor = 1_000_000;

        try galaxies1.append(Galaxy{ .row = row + expandedRows, .col = col + expandedColumns });
        try galaxies2.append(Galaxy{ .row = row + (expansionFactor - 1) * expandedRows, .col = col + (expansionFactor - 1) * expandedColumns });
    }

    // add up all their distances
    var part1: usize = 0;
    for (0..galaxies1.items.len - 1) |i| {
        for (i + 1..galaxies1.items.len) |j| {
            part1 += dist(galaxies1.items[j], galaxies1.items[i]);
        }
    }

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

    var part2: usize = 0;
    for (0..galaxies2.items.len - 1) |i| {
        for (i + 1..galaxies2.items.len) |j| {
            part2 += dist(galaxies2.items[j], galaxies2.items[i]);
        }
    }

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

fn dist(a: Galaxy, b: Galaxy) usize {
    const x = if (a.row > b.row) a.row - b.row else b.row - a.row;
    const y = if (a.col > b.col) a.col - b.col else b.col - a.col;
    return x + y;
}

fn findEmptyColumns(grid: std.ArrayList(std.ArrayList(u8)), allocator: std.mem.Allocator) !std.ArrayList(usize) {
    var cols = std.ArrayList(usize).init(allocator);
    const height = grid.items.len;
    const width = grid.items[0].items.len;

    for (0..width) |col| {
        var allEmpty = true;
        for (0..height) |row| {
            if (grid.items[row].items[col] == '#') {
                allEmpty = false;
            }
        }
        if (allEmpty) {
            try cols.append(col);
        }
    }
    return cols;
}

fn findEmptyRows(grid: std.ArrayList(std.ArrayList(u8)), allocator: std.mem.Allocator) !std.ArrayList(usize) {
    var rows = std.ArrayList(usize).init(allocator);

    for (grid.items, 0..) |row, idx| {
        var allEmpty = true;
        for (row.items) |char| {
            if (char == '#') {
                allEmpty = false;
            }
        }
        if (allEmpty) {
            try rows.append(idx);
        }
    }
    return rows;
}

fn parseGrid(content: []const u8, allocator: std.mem.Allocator) !std.ArrayList(std.ArrayList(u8)) {
    var grid = std.ArrayList(std.ArrayList(u8)).init(allocator);
    var lines = std.mem.split(u8, content, "\n");

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

        var al = std.ArrayList(u8).init(allocator);
        for (line) |char| {
            try al.append(char);
        }

        try grid.append(al);
    }

    return grid;
}