from collections import defaultdict
from dataclasses import dataclass
from functools import reduce
from typing import Any, Optional
from .aoc import Aoc2, AocSameParser


@dataclass
class Number:
    value: int
    symbol: Optional[tuple[int, int, str]]
    startx: int
    endx: int
    linenr: int


class Day03(AocSameParser[list[Number]], Aoc2[list[Number], list[Number]]):
    def __init__(self) -> None:
        super().__init__(2023, 3)

    def get_symbol(
        self, inpt: list[str], startx: int, endx: int, liney: int
    ) -> Optional[tuple[int, int, str]]:
        symbol = None
        for x in range(max(0, startx - 1), min(len(inpt[0]), endx + 1)):
            yrange = [liney - 1, liney, liney + 1]
            if x not in (endx, startx - 1):
                yrange.remove(liney)
            if liney == 0:
                yrange.remove(liney - 1)
            if liney == len(inpt) - 1:
                yrange.remove(liney + 1)
            for y in yrange:
                try:
                    if inpt[y][x] != "." and not inpt[y][x].isnumeric():
                        symbol = (x, y, inpt[y][x])
                except IndexError:
                    pass
        return symbol

    def parseinput(self, inpt: str) -> list[Number]:
        numbers = []
        lines = inpt.splitlines()
        for posy, line in enumerate(lines):
            num_accum = None
            firstx = 0
            lastx = 0
            for posx, char in enumerate(line):
                if char.isnumeric():
                    if num_accum is None:
                        num_accum = char
                        firstx = posx
                    else:
                        num_accum += char
                else:
                    if num_accum is not None:
                        lastx = posx
                        symbol = self.get_symbol(lines, firstx, lastx, posy)
                        numbers.append(
                            Number(int(num_accum), symbol, firstx, lastx, posy)
                        )

                        num_accum = None
                        firstx = 0
                        lastx = 0

            if num_accum is not None:
                lastx = len(line)
                symbol = self.get_symbol(lines, firstx, lastx, posy)
                numbers.append(Number(int(num_accum), symbol, firstx, lastx, posy))

                num_accum = None
                firstx = 0
                lastx = 0

        return numbers

    def part1(self, inpt: list[Number]) -> int:
        return sum(number.value for number in inpt if number.symbol is not None)

    def part2(self, inpt: list[Number]) -> Any:
        gears_to_number = defaultdict(list)
        for number in inpt:
            if number.symbol is not None and number.symbol[2] == "*":
                gears_to_number[number.symbol].append(number.value)
        return sum(
            reduce(lambda x, y: x * y, numbers)
            for numbers in gears_to_number.values()
            if len(numbers) == 2
        )


def main() -> None:
    day03 = Day03()
    day03.run()


if __name__ == "__main__":
    main()