import sys
from enum import Enum
from tools.aoc import AOCDay
from typing import Any


class OpCode(int, Enum):
    ADV = 0
    BXL = 1
    BST = 2
    JNZ = 3
    BXC = 4
    OUT = 5
    BDV = 6
    CDV = 7


class Computer:
    def __init__(self, reg_a: int, reg_b: int, reg_c: int, code: list[int]) -> None:
        self.__reg_a = reg_a
        self.__reg_b = reg_b
        self.__reg_c = reg_c
        self.reg_a = self.__reg_a
        self.reg_b = self.__reg_b
        self.reg_c = self.__reg_c
        self.output = []
        self.instr_pointer = 0
        self.__code = code
        self.code = [(OpCode(code[x]), code[x + 1]) for x in range(0, len(code), 2)]
        self.check_output = False
        self.op = {
            OpCode.ADV: self.adv,
            OpCode.BXL: self.bxl,
            OpCode.BST: self.bst,
            OpCode.JNZ: self.jnz,
            OpCode.BXC: self.bxc,
            OpCode.OUT: self.out,
            OpCode.BDV: self.bdv,
            OpCode.CDV: self.cdv,
        }

    def reset(self) -> None:
        self.reg_a = self.__reg_a
        self.reg_b = self.__reg_b
        self.reg_c = self.__reg_c
        self.output = []
        self.instr_pointer = 0

    def get_combo_value(self, value: int) -> int:
        if value < 0 or value > 6:
            raise ValueError(f"Invalid Combo Operand: {value}")

        if value <= 3:
            return value
        elif value == 4:
            return self.reg_a
        elif value == 5:
            return self.reg_b
        elif value == 6:
            return self.reg_c

    def adv(self, operand: int) -> None:
        self.reg_a = self.reg_a // (2 ** self.get_combo_value(operand))
        self.instr_pointer += 1

    def bxl(self, operand: int) -> None:
        self.reg_b ^= operand
        self.instr_pointer += 1

    def bst(self, operand: int) -> None:
        self.reg_b = self.get_combo_value(operand) % 8
        self.instr_pointer += 1

    def jnz(self, operand: int) -> None:
        if self.reg_a == 0:
            self.instr_pointer += 1
        else:
            if operand % 2 != 0:
                raise ValueError(f"Invalid JNZ Operand: {operand}")

            self.instr_pointer = operand // 2

    def bxc(self, _: int) -> None:
        self.reg_b = self.reg_c ^ self.reg_b
        self.instr_pointer += 1

    def out(self, operand: int) -> None:
        value = self.get_combo_value(operand) % 8
        self.output.append(value)
        if not self.check_output:
            self.instr_pointer += 1
        elif value != self.__code[len(self.output) - 1]:
            # print(f"Want to add {value} to {self.output=}, but expected {self.__code[len(self.output) - 1]}")
            self.instr_pointer = len(self.code)
        else:
            self.instr_pointer += 1

    def bdv(self, operand: int) -> None:
        self.reg_b = self.reg_a // (2 ** self.get_combo_value(operand))
        self.instr_pointer += 1

    def cdv(self, operand: int) -> None:
        self.reg_c = self.reg_a // (2 ** self.get_combo_value(operand))
        self.instr_pointer += 1

    def get_output(self) -> str:
        return ",".join(str(x) for x in self.output)

    def run_code(self) -> None:
        while self.instr_pointer < len(self.code):
            # print("Exec", self.instr_pointer)
            op_code, operand = self.code[self.instr_pointer]
            self.op[op_code](operand)


class Day(AOCDay):
    inputs = [
        [
            ("4,6,3,5,6,3,5,2,1,0", "input17_test"),
            ("1,6,3,6,5,6,5,1,7", "input17"),
        ],
        [
            # (117440, "input17_test2"),
            (None, "input17"),
        ],
    ]

    def parse_input(self) -> tuple[int, int, int, list[int]]:
        a, b, c, _, prog = self.getIntsFromInput()
        return a[0], b[0], c[0], prog

    def part1(self) -> Any:
        computer = Computer(*self.parse_input())
        computer.run_code()
        return computer.get_output()

    def part2(self) -> Any:
        a, b, c, prog = self.parse_input()
        expected_output = ",".join(str(x) for x in prog)
        computer = Computer(a, b, c, prog)
        print(computer.code)
        computer.check_output = True
        reg_a_init = 0
        ex_len = 1
        output = ""
        while output != expected_output:
            reg_a_init += 1
            computer.reset()
            computer.reg_a = reg_a_init
            computer.run_code()
            output = computer.get_output()
            if output.startswith(expected_output[:ex_len]):
                print(f"Computer Reg A: {reg_a_init}, BXL: {reg_a_init ^ 1}")
                print(f"Computer Output: {output}")
                sys.stdout.flush()
                ex_len += 2

        return reg_a_init


if __name__ == "__main__":
    day = Day(2024, 17)
    day.run(verbose=True)