Advent of Code 2024 day 17 part 2

from sys import argv


def parse_register(line):
    index = line.find(':') + 2
    end_index = index + 1
    while end_index <= len(line) and line[index:end_index].isdigit():
        end_index += 1

    return int(line[index:end_index])


def parse_program(line):
    index = line.find(':') + 2
    return [int(x) for x in line[index:].split(',')]


def parse(src):
    _, program = src.split('\n\n', 1)
    program = parse_program(program)
    return program


def get_register(a, b, c, opcode):
    if opcode <= 3:  # assume the opcode is in bounds
        return opcode
    elif opcode == 4:
        return a
    elif opcode == 5:
        return b
    else:
        return c


def test(a, program, end, out_index):
    b = c = 0
    pc = 0
    while pc < end:
        instruction, opcode = program[pc:pc+2]
        match instruction:
            case 0:
                a = a >> get_register(a, b, c, opcode)

            case 1:
                b ^= opcode

            case 2:
                b = get_register(a, b, c, opcode) % 8

            case 3:
                if a != 0:
                    pc = opcode
                    continue

            case 4:
                b ^= c

            case 5:
                output = get_register(a, b, c, opcode) % 8
                if output != program[out_index]:
                    return False
                out_index += 1

            case 6:
                b = a >> get_register(a, b, c, opcode)

            case 7:
                c = a >> get_register(a, b, c, opcode)

        pc += 2

    return out_index == end


def main(program):
    end = len(program)

    # 0 gets the initial 0 and no more
    in_progress = [(0, len(program) - 1)]
    while len(in_progress) > 0:
        a, distance = in_progress.pop(0)
        for chunk in range(8):
            next_a = (a << 3) + chunk

            if not test(next_a, program, end, distance):
                continue

            if distance == 0:
                return next_a

            in_progress.append((next_a, distance - 1))

    return None


if __name__ == '__main__':
    print(main(parse(open(argv[1]).read())))