const std = @import("std");

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

pub fn main() !void {
    const grid = @embedFile("data.txt");
    const width = std.mem.indexOf(u8, grid, "\n").?; // can't do this at comptime. damn
    const height = grid.len / width - 1;
    const start = std.mem.indexOf(u8, grid, "S").?;
    const startCoord = indexToCoord(start, width);
    const allStartingNeighbors = neighbors(startCoord, width, height);

    // find the 2 path starting points
    var paths: [2]Coord = undefined;
    var pathIndex: u2 = 0;

    for (allStartingNeighbors, 0..) |neighbor, dir| {
        if (neighbor) |c| {
            const value = grid[coordToIndex(c, width)];
            var validPath = false;
            if (dir == 0 and (value == '|' or value == 'F' or value == '7')) {
                validPath = true;
            }
            if (dir == 1 and (value == '-' or value == 'J' or value == '7')) {
                validPath = true;
            }
            if (dir == 2 and (value == '|' or value == 'J' or value == 'L')) {
                validPath = true;
            }
            if (dir == 3 and (value == '-' or value == 'F' or value == 'L')) {
                validPath = true;
            }
            if (validPath) {
                paths[pathIndex] = c;
                pathIndex += 1;
            }
        }
    }

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

    // capture the loop's coordinates in a map
    var loop = std.ArrayList(Coord).init(allocator);
    defer loop.deinit();
    try loop.append(startCoord);
    try loop.append(paths[0]);

    // count the steps
    var steps: u16 = 1;
    var totalSteps: u16 = 0;
    var prevs = [2]Coord{ startCoord, startCoord };

    while (true) {
        // for each of the 2 paths
        for (0..2) |i| {
            const temp = visit(paths[i], prevs[i], grid, width).?;

            if (i == 0) {
                // try loop.put(temp, {});
                try loop.append(temp);
            }

            prevs[i] = paths[i];
            paths[i] = temp;
        }

        steps += 1;

        // if the two paths converge (and not at the starting point again)
        if (paths[0].row == paths[1].row and paths[0].col == paths[1].col and (paths[0].row != startCoord.row or paths[0].col != startCoord.col)) {
            totalSteps = steps;
        }

        // for part 2, continue the first path until it gets back to the beginning
        if (paths[0].row == startCoord.row and paths[0].col == startCoord.col) {
            break;
        }
    }

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

    // Part 2
    var visited = std.AutoHashMap(Coord, void).init(allocator);
    defer visited.deinit();

    var prev = startCoord;
    for (loop.items) |c| {
        // find the coord to the "right"
        var coord: ?Coord = null;
        if (prev.row < c.row) {
            coord = Coord{ .row = c.row, .col = c.col - 1 };
        } else if (prev.row > c.row) {
            coord = Coord{ .row = c.row, .col = c.col + 1 };
        } else if (prev.col < c.col) {
            coord = Coord{ .row = c.row + 1, .col = c.col };
        } else {
            coord = Coord{ .row = c.row - 1, .col = c.col };
        }

        try visitAllNeighbors(coord.?, &visited, &loop, grid, width, height);
        prev = c;
    }

    // looking at my printed grid, there are a bunch of stragglers. Donno how. Let's count them
    // var stragglers = std.ArrayList(Coord).init(allocator);
    // defer stragglers.deinit();
    //
    // for (grid) |i| {
    //     const c = indexToCoord(i, width);
    //     const ns = neighbors(c, width, height);
    //     var allVisited = true;
    //     for (ns) |n| {
    //         if (n != null and visited.get(n.?) != null) {
    //             allVisited = false;
    //         }
    //     }
    //     if (allVisited) {
    //         try visited.put(c, {});
    //     }
    // }

    printGrid(grid, width, visited, &loop);
    std.debug.print("Part 2: {any}\n", .{visited.count()});
}

fn printGrid(grid: []const u8, width: usize, visited: anytype, loop: anytype) void {
    for (grid, 0..) |value, idx| {
        const c = indexToCoord(idx, width);
        if (contains(loop, c)) {
            // std.debug.print("{c}", .{value});
            std.debug.print(" ", .{});
        } else if (visited.get(c) != null) {
            std.debug.print("X", .{});
        } else if (value == '\n') {
            std.debug.print("{c}", .{value});
        } else {
            // std.debug.print("{c}", .{value});
            std.debug.print(".", .{});
        }
    }
}

fn visitAllNeighbors(c: Coord, visited: anytype, loop: anytype, grid: []const u8, width: usize, height: usize) !void {
    // if it's been visited, get out
    if (visited.get(c) != null) {
        return;
    }

    if (contains(loop, c)) {
        return;
    }

    try visited.put(c, {});

    const ns = neighbors(c, width, height);
    for (ns) |neighbor| {
        if (neighbor) |n| {
            try visitAllNeighbors(n, visited, loop, grid, width, height);
        }
    }
}

fn contains(ar: *std.ArrayList(Coord), item: Coord) bool {
    for (ar.items) |i| {
        if (i.row == item.row and i.col == item.col) {
            return true;
        }
    }

    return false;
}

// pass it a coordinate and the previous coordinate, and it finds the next one
fn visit(c: Coord, prev: Coord, grid: []const u8, width: usize) ?Coord {
    const value = grid[coordToIndex(c, width)];

    return switch (value) {
        '-' => Coord{ .row = c.row, .col = if (prev.col < c.col) c.col + 1 else c.col - 1 },
        '|' => Coord{ .row = if (prev.row < c.row) c.row + 1 else c.row - 1, .col = c.col },
        'F' => if (prev.row > c.row) Coord{ .row = c.row, .col = c.col + 1 } else Coord{ .row = c.row + 1, .col = c.col },
        '7' => if (prev.row > c.row) Coord{ .row = c.row, .col = c.col - 1 } else Coord{ .row = c.row + 1, .col = c.col },
        'J' => if (prev.row < c.row) Coord{ .row = c.row, .col = c.col - 1 } else Coord{ .row = c.row - 1, .col = c.col },
        'L' => if (prev.row < c.row) Coord{ .row = c.row, .col = c.col + 1 } else Coord{ .row = c.row - 1, .col = c.col },
        else => null,
    };
}

// returns all neighboring coordinates, or null if it would be off the grid
fn neighbors(c: Coord, width: usize, height: usize) [4]?Coord {
    const up = if (c.row > 0) Coord{ .row = c.row - 1, .col = c.col } else null;
    const right = if (c.col < width - 1) Coord{ .row = c.row, .col = c.col + 1 } else null;
    const down = if (c.row < height - 1) Coord{ .row = c.row + 1, .col = c.col } else null;
    const left = if (c.col > 0) Coord{ .row = c.row, .col = c.col - 1 } else null;

    return [_]?Coord{
        up,
        right,
        down,
        left,
    };
}

fn coordToIndex(c: Coord, width: usize) usize {
    return c.row * (width + 1) + c.col;
}

fn indexToCoord(index: usize, width: usize) Coord {
    return Coord{ .row = index / (width + 1), .col = index % (width + 1) };
}