import re
from typing import override
from unittest import TestCase

from puzzles._solver import Solver

simple_instruction_pattern = re.compile(r"mul\((\d+),(\d+)\)")
instruction_pattern = re.compile(r"mul\((\d+),(\d+)\)|(do(?:n't)?)\(\)")


class DayNSolver(Solver):
    memory: str

    @override
    def __init__(self, puzzle_input: str):
        self.memory = puzzle_input

    @override
    def solve_p1(self) -> int:
        total = 0

        for match in simple_instruction_pattern.findall(self.memory):
            a, b = match
            total += int(a) * int(b)

        return total

    @override
    def solve_p2(self) -> int:
        total = 0
        enabled = True

        for match in instruction_pattern.findall(self.memory):
            a, b, enable_flag = match
            match enable_flag:
                case "do":
                    enabled = True
                case "don't":
                    enabled = False
                case "":
                    if enabled:
                        total += int(a) * int(b)

        return total


class TestDayNSolver(TestCase):
    def test_solve_p1(self):
        solver = DayNSolver(
            "xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))\n"
        )
        self.assertEqual(solver.solve_p1(), 161)

    def test_solve_p2(self):
        solver = DayNSolver(
            "xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))\n"
        )
        self.assertEqual(solver.solve_p2(), 48)