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.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)
        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(self, init_a: int = None) -> str:
        self.reset()
        if init_a is not None:
            self.reg_a = init_a

        while self.instr_pointer < len(self.code):
            op_code, operand = self.code[self.instr_pointer]
            self.op[op_code](operand)

        return self.get_output()


class Day(AOCDay):
    inputs = [
        [
            ("4,6,3,5,6,3,5,2,1,0", "input17_test"),
            ("7,1,2,3,2,6,7,2,5", "input17_neil"),
            ("1,7,6,5,1,0,5,0,7", "input17_dennis"),
            ("1,6,3,6,5,6,5,1,7", "input17"),
        ],
        [
            (117440, "input17_test2"),
            (202356708354602, "input17_neil"),
            (236555995274861, "input17_dennis"),
            (247839653009594, "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())
        return computer.run()

    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)

        def get_match_pos(a: str, b: str) -> int:
            m = 0
            for x, y in zip(reversed(a), reversed(b)):
                if x == y:
                    m += 1
                else:
                    break

            return m

        num = 0
        for factor in range(len(prog)):
            for i in range(8**len(prog)):
                output = computer.run(num + i)
                if get_match_pos(output, expected_output) == (2 * factor) + 1:
                    if factor < len(prog) - 1:
                        num = (i + num) * 8
                    else:
                        num += i
                    break

        return num


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