from typing import Iterator
import sys
import re

mul_re = re.compile(
    r"mul\((?P<n1>\d+),(?P<n2>\d+)\)"
)

mul_re_cond = re.compile(
    r"mul\((?P<n1>\d+),(?P<n2>\d+)\)|(?P<do>do\(\))|(?P<dont>don't\(\))"
)

def process_basic(data) -> int:
    total = 0
    for match in re.finditer(mul_re, data):
        n1 = int(match.group('n1'))
        n2 = int(match.group('n2'))
        total += n1 * n2
    return total

def process_cond(data):
    # Could also split into segments between do and the next don't
    # And process each block between them
    # But this is simpler
    total = 0
    enabled = True
    for match in re.finditer(mul_re_cond, data):
        if match.group('do'):
            enabled = True
        elif match.group('dont'):
            enabled = False
        elif enabled:
            n1 = int(match.group('n1'))
            n2 = int(match.group('n2'))
            total += n1 * n2
    return total


def load_data(filename) -> str:
    with open(filename) as f:
        return f.read()

def test():
    print("Beginning Test 1")
    data = r"xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))"
    assert process_basic(data) == 161
    print("Test 1 passed")

    print("Beginning Test 2")
    data = r"xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))"
    assert process_cond(data) == 48
    print("Test 2 passed")

def main():
    test()
    data = load_data(sys.argv[1])
    total = process_basic(data)
    print(f"Total from data without conditionals: {total}")
    total = process_cond(data)
    print(f"Total from data with conditionals: {total}")





if __name__ == '__main__':
    main()