#!/usr/bin/env python3

# https://adventofcode.com/2024/day/6

def get_lines(filename: str) -> list:
    with open(filename, "r") as file:
        return [line.strip() for line in file.readlines()]


def get_lab(lines: list) -> tuple:
    lab = []
    g_row = 0
    g_col = 0
    for row_num, row_content in enumerate(lines):
        row = []
        for col_num, cell_content in enumerate(row_content):
            if cell_content == "^":
                g_row, g_col = row_num, col_num
                row.append('.')
            else:
                row.append(cell_content)
        lab.append(row)
    return lab, g_row, g_col


def next_position(row: int, col: int, direction: int) -> tuple:
    deltas = [(-1, 0), (0, 1), (1, 0), (0, -1)]
    delta_row, delta_col = deltas[direction]
    new_row = row + delta_row
    new_col = col + delta_col
    return new_row, new_col


def pos_invalid(row, col, max_row, max_col) -> bool:
    if 0 <= row <= max_row and 0 <= col <= max_col:
        return False
    return True


def find_loop(lab: list, g_row: int, g_col: int, g_dir: int,
              block_row: int, block_col: int) -> bool:
    max_row = len(lab) - 1
    max_col = len(lab[0]) - 1
    if pos_invalid(block_row, block_col, max_row, max_col):
        return False
    test_visits = set()
    while True:
        if (g_row, g_col, g_dir) in test_visits:
            return True
        test_visits.add((g_row, g_col,g_dir))
        new_row, new_col = next_position(g_row, g_col, g_dir)
        if pos_invalid(new_row, new_col, max_row, max_col):
            return False
        elif lab[new_row][new_col] == "#" or (new_row == block_row and new_col == block_col):
            g_dir = (g_dir + 1) % 4
        else:
            g_row = new_row
            g_col = new_col


def get_visits(lab: list, g_row: int, g_col: int) -> tuple:
    visits = set()
    obstructions = set()
    obs_tested = set()
    g_dir = 0
    start_row, start_col, start_dir = g_row, g_col, g_dir
    max_row = len(lab) - 1
    max_col = len(lab[0]) - 1
    while True:
        visits.add((g_row, g_col))
        new_row, new_col = next_position(g_row, g_col, g_dir)
        if pos_invalid(new_row, new_col, max_row, max_col):
            return len(visits), len(obstructions)
        elif lab[new_row][new_col] == "#":
            g_dir = (g_dir + 1) % 4
        else:
            block_row, block_col = next_position(g_row, g_col, g_dir)
            if not (block_row, block_col) in obs_tested:
                if find_loop(lab, start_row, start_col, start_dir, block_row, block_col):
                    obstructions.add((block_row,block_col))
                obs_tested.add((block_row, block_col))
            g_row = new_row
            g_col = new_col


def main():
    # lines = get_lines("sample-input.txt")
    lines = get_lines("input.txt")
    lab, g_row, g_col = get_lab(lines)
    visits, obstructions = get_visits(lab, g_row, g_col)
    print("Part 1: Visited positions:", visits)
    print("Part 2: Possible positions for obstructions:", obstructions)


if __name__ == '__main__':
    main()