import re
from lib import *

EXAMPLE = """
467..114..
...*......
..35..633.
......#...
617*......
.....+.58.
..592.....
......755.
...$.*....
.664.598..
"""

def solve(i: Input, second=False):
    """
    This is a new version I have implemented after the fact to test the
    improvements to my library. My original versions are `solve1` and `solve2`.
    """
    g = i.grid2()
    if not second:
        parts = g.find_not(NUMBERS + ".")
    else:
        parts = g.find("*")

    res = 0
    numbers = []
    for p in parts:
        numbers_gear = []
        for n in g.neighbors_adj(p):
            if g[n] in NUMBERS:
                while n in g and g[n] in NUMBERS:
                    n = (n[0], n[1] - 1)
                number = ""
                n = (n[0], n[1] + 1)
                start = n
                while n in g and g[n] in NUMBERS:
                    number += g[n]
                    n = (n[0], n[1] + 1)
                numbers.append((int(number), n[0], start))
                numbers_gear.append(int(number))
        numbers_gear = list(set(numbers_gear))
        if len(numbers_gear) == 2:
            res += numbers_gear[0] * numbers_gear[1]

    if second:
        return res
    else:
        return sum([n for n, _, _ in list(set(numbers))])

def clean(text: str) -> list[str]:
    return list(filter(lambda l: l.strip() != "", text.splitlines()))

def is_adj_to_symbol(x, y, lines):
    non_symbol = re.compile(r"[^0-9\.]")
    for xo, yo in [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]:
        try:
            if x + xo < 0 or y + yo < 0:
                continue
            if non_symbol.match(lines[y + yo][x + xo]):
                return True
        except IndexError:
            pass
    return False

def solve1(lines: list[str]):
    d = ""
    adj_to_symbol = False
    s = 0
    for (y, line) in enumerate(lines):
        for (x, c) in enumerate(line):
            if c not in NUMBERS:
                if len(d) > 0 and adj_to_symbol:
                    s += int(d)
                d = ""
                adj_to_symbol = False
            else:
                if is_adj_to_symbol(x, y, lines):
                    adj_to_symbol = True
                d = d + c
    return s

def get_entire_number(x, y, lines):
    assert(lines[y][x] in NUMBERS)
    x_start = x - 1
    while lines[y][x_start] in NUMBERS and x_start >= 0:
        x_start -= 1
    x_start += 1
    x_start_orig = x_start
    r = ""
    while x_start < len(lines[0]) and lines[y][x_start] in NUMBERS:
        r += lines[y][x_start]
        x_start += 1
    return x_start_orig, int(r)

def find_sourrounding_numbers(x, y, lines):
    full_coords = []
    numbers = []
    for xo, yo in [(-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1)]:
        new_x = x + xo
        new_y = y + yo
        if new_x < 0 or new_y < 0 or new_x >= len(lines[0]) or new_y >= len(lines):
            continue
        if lines[new_y][new_x] in NUMBERS:
            x_start_orig, n = get_entire_number(new_x, new_y, lines)
            if (x_start_orig, new_y, n) not in full_coords:
                numbers.append(n)
                full_coords.append((x_start_orig, y + yo, n))
    return numbers

def solve2(lines: list[str]):
    s = 0
    for (y, line) in enumerate(lines):
        for (x, c) in enumerate(line):
            if c == '*':
                numbers = find_sourrounding_numbers(x, y, lines)
                if len(numbers) == 2:
                    # print(numbers)
                    s += numbers[0] * numbers[1]
    return s

def main():
    input = Input(EXAMPLE)
    print("Example  1:", solve(input))

    input = Input("i3.txt")
    print("Solution 1:", solve(input))

    input = Input(EXAMPLE)
    print("Example  2:", solve(input, True))

    input = Input("i3.txt")
    print("Solution 2:", solve(input, True))
    return

if __name__ == "__main__":
    main()