import re
from itertools import combinations
from math import ceil
from tools.aoc import AOCDay
from typing import Any


re_pair = re.compile(r"\[(\d+),(\d+)\]")
re_lnum = re.compile(r".*[^\d](\d+)")
re_rnum = re.compile(r"(\d+)")
re_bignum = re.compile(r"(\d\d+)")


class Snailfish:
    def __init__(self, fish_str: str):
        self.pairs = fish_str
        while self.reduce():
            pass

    def __add__(self, other: 'Snailfish'):
        return Snailfish("[%s,%s]" % (self.pairs, other.pairs))

    def reduce(self):
        open_count = 0
        for i, c in enumerate(self.pairs):
            if c == "]":
                open_count -= 1
            elif c == "[":
                open_count += 1
                if open_count == 5:
                    explode_pair_match = re_pair.search(self.pairs, i)
                    index_l, index_r = explode_pair_match.span()
                    explode_left, explode_right = map(int, explode_pair_match.groups())
                    lnum_old = rnum_old = lnum_new = rnum_new = ""
                    lnum_index = index_l
                    rnum_index = index_r

                    if left_num := re_lnum.search(self.pairs[:index_l]):
                        lnum_old = left_num.groups()[0]
                        lnum_new = str(int(lnum_old) + explode_left)
                        lnum_index = left_num.span()[1] - len(lnum_old)

                    if right_num := re_rnum.search(self.pairs, index_r):
                        rnum_old = right_num.groups()[0]
                        rnum_new = str(int(rnum_old) + explode_right)
                        rnum_index = right_num.span()[0]

                    self.pairs = "".join([
                        self.pairs[:lnum_index],
                        lnum_new,
                        self.pairs[lnum_index + len(lnum_old):index_l],
                        "0",
                        self.pairs[index_r:rnum_index],
                        rnum_new,
                        self.pairs[rnum_index + len(rnum_old):]
                    ])
                    return True

        if big_num_match := re_bignum.search(self.pairs):
            num = int(big_num_match.group())
            self.pairs = self.pairs.replace(big_num_match.group(), "[%d,%d]" % (num // 2, int(ceil(num / 2))), 1)
            return True

        return False

    def getMagnitude(self) -> int:
        pairs = self.pairs
        while match := re_pair.findall(pairs):
            for num1, num2 in match:
                pairs = pairs.replace("[%s,%s]" % (num1, num2), str(int(num1) * 3 + int(num2) * 2))

        return int(pairs)


class Day(AOCDay):
    test_solutions_p1 = [4140, 4417]
    test_solutions_p2 = [3993, 4796]

    def part1(self) -> Any:
        snailfishes = [Snailfish(x) for x in self.getInput()]
        return sum(snailfishes[1:], start=snailfishes[0]).getMagnitude()

    def part2(self) -> Any:
        snailfishes = [Snailfish(x) for x in self.getInput()]
        max_mag = 0
        for a, b in combinations(snailfishes, 2):
            sub_mag_a = (a + b).getMagnitude()
            sub_mag_b = (b + a).getMagnitude()
            max_mag = max(max_mag, sub_mag_a, sub_mag_b)

        return max_mag