with open('input.text', 'r') as file:
    data: list[str] = file.readlines()

y_bound = len(data)
x_bound = len(data[0])

# Part 1
class Vector:
    x: int
    y: int

    def __init__(self, x: int, y: int):
        self.x = x
        self.y = y

    def __add__(self, other):
        return Vector(self.x + other.x, self.y + other.y)

    def __str__(self):
        return f'{self.x}, {self.y}'

DOWN = Vector(0, -1)
UP = Vector(0, 1)
RIGHT = Vector(1, 0)
LEFT = Vector(-1, 0)

debug_grid = ['.' * x_bound for n in range(y_bound)]

def search_direction(pos: Vector, dir: Vector) -> bool:
    poss = []
    s = data[pos.y][pos.x]

    poss.append(pos)
    for i in range(0, 3):
        pos = pos + dir
        poss.append(pos)

        # Check Y
        if pos.y < 0 or pos.y >= len(data):
            return False

        # Check X
        elif pos.x < 0 or pos.x >= len(data[pos.y]):
            return False

        s += data[pos.y][pos.x]

        if s != 'X' and s != 'XM' and s != 'XMA' and s != 'XMAS':
            return False

    for i in range(0, 4):
        debug_grid[poss[i].y] = debug_grid[poss[i].y][:poss[i].x] + s[i] + debug_grid[poss[i].y][poss[i].x + 1:]

    return True

total = 0
for y, line in enumerate(data):
    for x, char in enumerate(line):
        if char == 'X':
            pos = Vector(x, y)
            total += (search_direction(pos, UP)
                      + search_direction(pos, DOWN)
                      + search_direction(pos, LEFT)
                      + search_direction(pos, RIGHT)
                      + search_direction(pos, UP + LEFT)
                      + search_direction(pos, UP + RIGHT)
                      + search_direction(pos, DOWN + LEFT)
                      + search_direction(pos, DOWN + RIGHT)
                )

for row in debug_grid:
    print(row)

print(f'XMAS Count: {total}')

# Part 2
def check_pos(pos: Vector) -> bool:
    chars = []
    poss = []

    for dir in [UP + LEFT, UP + RIGHT, DOWN + LEFT, DOWN + RIGHT]:
        new_pos = pos + dir

        # Check Y
        if new_pos.y < 0 or new_pos.y >= len(data):
            return False

        # Check X
        elif new_pos.x < 0 or new_pos.x >= len(data[new_pos.y]):
            return False
        chars.append(data[new_pos.y][new_pos.x])
        poss.append(new_pos)

    if chars.count('M') == 2 and chars.count('S') == 2:
        m_pos = None

        for i, char in enumerate(chars):
            match char:
                case 'M':
                    if not m_pos:
                        m_pos = poss[i]
                    else:
                        if poss[i].x != m_pos.x and poss[i].y != m_pos.y:
                            return False



        return True

    return False


total = 0
for y, line in enumerate(data):
    for x, char in enumerate(line):
        if char == 'A':
            total += check_pos(Vector(x, y))

print(f'X-MAS Count: {total}')