import re
from typing import Literal, override
from unittest import TestCase

from puzzles._solver import Solver


class DayThreeSolver(Solver):
    memory: tuple[int, int] | Literal["do", "don't"]

    @override
    def __init__(self, puzzle_input: str):
        matches = re.findall(r"mul\((\d+),(\d+)\)|(do(?:n't)?)\(\)", puzzle_input)
        self.memory = tuple(
            (int(a), int(b)) if a else enable_flag for a, b, enable_flag in matches
        )

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

        for match in self.memory:
            match match:
                case a, b:
                    total += a * b

        return total

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

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

        return total


class TestDayThreeSolver(TestCase):
    def test_solve_p1(self):
        solver = DayThreeSolver(
            "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 = DayThreeSolver(
            "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)