from typing import Tuple, List, Dict, Set
from dataclasses import dataclass
from collections import defaultdict

@dataclass(frozen=True)
class Item:
    pos: Tuple[int, int]
    symbol: str

def browse_schema(schema):
    total_parts = 0
    buf = []
    max_row, max_col = len(schema), len(schema[0])

    numbers: List[Item] = []

    for y in range(max_row):
        for x in range(max_col):
            item = schema[y][x]
            if item.isnumeric():
                # continue parsing full number
                buf.append(item)
            elif buf:
                # end of a number, do the engine part check
                number = "".join(buf)
                neighbors = get_neighbors((x, y), len(buf), schema)
                start_pos = (x-len(number), y)
                numbers.append(Item(start_pos, number))
                
                if is_engine_part(neighbors):
                    total_parts += int(number)
                
                buf.clear()  # reached end of a number, clear buffer

    print(f"Part 1, sum of the parts numbers = {total_parts}")

    gears: Dict[Tuple[int, int], Set[Item]] = defaultdict(set)
    for y in range(max_row):
        for x in range(max_col):
            item = schema[y][x]
            if item == "*":
                print(f"* ({x}, {y})")
                neighbors = [n for n in get_neighbors_of((x, y), schema) if n.symbol.isdigit()]
                if (x, y) == (5, 8):
                    breakpoint()

                for neighbor in neighbors:
                    neighboring_numbers = [n for n in numbers
                        if y-1 <= n.pos[1] == y+1
                        and n.pos[0] <= neighbor.pos[0] <= n.pos[0] + len(n.symbol)]
                    gears[(x,y)] = gears[(x,y)].union(neighboring_numbers)
    
    # gears = { key: value for key, value in gears.items() if len(value) <= 2 }



    print(f"Part 2, sum of the gear ratios = TODO")


def is_engine_part(neighbors: List[Item]) -> bool:
    # get list of symbols (not '.', \n or a number)
    symbols = filter(lambda x: not x.symbol.isnumeric() and not x.symbol in (".", "\n"), neighbors)
    return next(symbols, None) is not None


def get_neighbors(pos: Tuple[int, int], length: int,  schema: List[List[str]]) -> List[Item]:
    x, y = pos
    start_x = x - length
    neighbors = [get_neighbors_of((x, y), schema) for x in range(start_x, x)]
    neighbors = [item for sublist in neighbors for item in sublist]  # flatten list of list
    return neighbors


def get_neighbors_of(pos: Tuple[int, int], schema: List[List[str]]) -> List[Item]:
    max_row, max_col = len(schema), len(schema[0])
    x, y = pos
    neighbors: List[Item] = []
    
    # top
    if y-1 >= 0:
        neighbors.append(Item((x, y-1), schema[y-1][x]))
    # bottom:
    if y+1 < max_row:
        neighbors.append(Item((x, y+1), schema[y+1][x]))
    # left
    if x-1 >= 0:
        neighbors.append(Item((x-1, y), schema[y][x-1]))
    # right
    if x+1 < max_col:
        neighbors.append(Item((x+1, y), schema[y][x+1]))
    # top-left
    if y-1 >= 0 and x-1 >= 0:
        neighbors.append(Item((x-1, y-1), schema[y-1][x-1]))
    # top-right
    if y-1 >= 0 and x+1 < max_col:
        neighbors.append(Item((x+1, y-1), schema[y-1][x+1]))
    # bottom-left
    if y+1 < max_row and x-1 >= 0:
        neighbors.append(Item((x-1, y+1), schema[y+1][x-1]))
    # bottom-right
    if y+1 < max_row and x+1 < max_col:
        neighbors.append(Item((x+1, y+1), schema[y+1][x+1]))

    return neighbors


if __name__ == "__main__":
    import sys
    infile = sys.argv[1]
    with open(infile) as f:
        schema = [[c for c in line] for line in f.readlines()]
        browse_schema(schema)