Advent of Code 2024 day 17 part 1

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):
    registers, program = src.split('\n\n', 1)

    registers = (parse_register(line) for line in registers.split('\n')
                 if line != '')

    program = parse_program(program)

    return registers, program, len(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 main(registers, program, end):
    a, b, c = registers
    pc = 0
    output = list()
    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.append(str(get_register(a, b, c, opcode) % 8))

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

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

        pc += 2

    return ','.join(output)


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