import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class Machine:
    register_a: int
    register_b: int
    register_c: int

    ic: int = 0 # Instruction Counter

    program: list[int] = []

    def __init__(self, program: str):
        self.register_a = 0
        self.register_b = 0
        self.register_c = 0

        self.program = [int(x) for x in program.split(',')]

    def interpret_combo_operand(self, operand: int) -> int:
        match operand:
            case operand if operand in range(0, 4):
                return operand
            case 4:
                return self.register_a
            case 5:
                return self.register_b
            case 6:
                return self.register_c

    def adv(self, operand: int) -> None:
        operand = self.interpret_combo_operand(operand)
        self.ic += 2
        self.register_a = int(self.register_a / (2 ** operand))

    def bxl(self, operand: int) -> None:
        self.ic += 2
        self.register_b = self.register_b ^ operand

    def bst(self, operand: int) -> None:
        operand = self.interpret_combo_operand(operand)
        self.ic += 2
        self.register_b = operand % 8

    def jnz(self, operand: int) -> None:
        if self.register_a == 0:
            self.ic += 2
            return
        self.ic = operand

    def bxc(self, operand: int) -> None:
        self.ic += 2
        self.register_b = self.register_b ^ self.register_c

    def out(self, operand: int) -> int:
        operand = self.interpret_combo_operand(operand)
        self.ic += 2
        return operand % 8

    def bdv(self, operand: int) -> None:
        operand = self.interpret_combo_operand(operand)
        self.ic += 2
        self.register_b = int(self.register_a / (2 ** operand))

    def cdv(self, operand: int) -> None:
        operand = self.interpret_combo_operand(operand)
        self.ic += 2
        self.register_c = int(self.register_a / (2 ** operand))

    def run(self) -> str:
        output = []
        while self.ic < len(self.program):

            opcode = self.program[self.ic]
            operand = self.program[self.ic + 1]
            logging.debug(f"IC: {self.ic}, A: {self.register_a}, B: {self.register_b}, C: {self.register_c}, Opcode: {opcode}, Operand: {operand}")

            opcode_out: None | int = None
            match opcode:
                case 0:
                    opcode_out = self.adv(operand)
                case 1:
                    opcode_out = self.bxl(operand)
                case 2:
                    opcode_out = self.bst(operand)
                case 3:
                    self.jnz(operand)
                case 4:
                    opcode_out = self.bxc(operand)
                case 5:
                    opcode_out = self.out(operand)
                case 6:
                    self.bdv(operand)
                case 7:
                    self.cdv(operand)

            if opcode_out is not None:
                output.append(opcode_out)

        return ','.join([str(x) for x in output])

def part1():
    INPUT_FILE = "input.txt"

    with open(INPUT_FILE, 'r') as file:
        data = file.readlines()

    reg_a = int(data[0].split(':')[1].strip())
    reg_b = int(data[1].split(':')[1].strip())
    reg_c = int(data[2].split(':')[1].strip())

    program = data[4].split(':')[1].strip()

    machine = Machine(program)
    machine.register_a = reg_a
    machine.register_b = reg_b
    machine.register_c = reg_c

    print(f'Output: {machine.run()}')
    logging.debug(f'Final Registers: A={machine.register_a}, B={machine.register_b}, C={machine.register_c}')

def part2():
    INPUT_FILE = "input.txt"

    with open(INPUT_FILE, 'r') as file:
        data = file.readlines()

    reg_b = int(data[1].split(':')[1].strip())
    reg_c = int(data[2].split(':')[1].strip())

    program = data[4].split(':')[1].strip()

    program_len: int = len(program.split(','))

    reg_a = 8 ** program_len - 1

    # machine = Machine(program)
    # machine.register_a = reg_a
    # machine.register_b = reg_b
    # machine.register_c = reg_c

    # print(machine.run())
    # print(program)

    # exit()

    # Naive solution but it works (in time)
    while True:
        logging.info(f'Testing Reg A: {reg_a}')
        machine = Machine(program)
        machine.register_a = reg_a
        machine.register_b = reg_b
        machine.register_c = reg_c

        output = machine.run()
        logging.info(f'Output: {output}')

        if output == program:
            print(f'Correct Reg A: {reg_a}')
            break

        reg_a += 1

if __name__ == "__main__":
    part2()