import math
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
        if not self.check_output:
            self.instr_pointer += 1
            self.output.append(value)
        elif value != self.__code[len(self.output)]:
            self.instr_pointer = len(self.code)
        else:
            self.instr_pointer += 1
            self.output.append(value)

    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"),
            ("7,1,2,3,2,6,7,2,5", "input17_neil"),
            ("1,6,3,6,5,6,5,1,7", "input17"),
        ],
        [
            (117440, "input17_test2"),
            (202356708354602, "input17_neil"),
            (236555995274861, "input17_dennis"),
            (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)
        computer.check_output = True

        def run_computer(v: int) -> str:
            computer.reset()
            computer.reg_a = v
            computer.run_code()
            return computer.get_output()

        init_a = 0
        init_reg = 0
        output = ""
        bit_match = ""
        bit_width = 6
        bit_add = bit_width - 1
        factor = 0
        additive = 0
        best = 0
        while output != expected_output:
            init_reg = init_a * (8**factor) + additive
            output = run_computer(init_reg)

            if len(output) > best and output == expected_output[: len(output)]:
                # if output == expected_output[: len(output)]:
                o = oct(init_reg)
                if o[-bit_width:] != bit_match:
                    bit_match = o[-bit_width:]
                else:
                    additive = int(bit_match, 8)
                    factor = bit_width
                    bit_width += bit_add
                    bit_add -= 1
                    init_a = 0

                print(
                    f"{init_reg=} ({oct(init_reg)}), {factor=}, additive={oct(additive)}, {output=}, {expected_output=}"
                )
                best = len(output)

            if init_a % 1_000_000 == 0:
                print(init_a, init_reg)
            init_a += 1

        print(init_a - 1, (init_a - 1) * (8**factor) + additive)
        output = run_computer(init_reg)
        print(output)
        if output == expected_output:
            return init_reg
        return ""

        first = []
        for i, c in enumerate(prog):
            reg_a_init = 0

            while len(first) <= i:
                computer.reset()
                computer.reg_a = reg_a_init
                computer.run_code()
                if int(computer.get_output()[0]) == c:
                    first.append(reg_a_init)

                reg_a_init += 1

        print(first)
        result = 0
        for i, c in enumerate(reversed(first)):
            result = result << 3
            result += c // 8
        # result = sum(((x * (8 ** i) for i, x in enumerate(first))))
        print(result)

        computer.reset()
        computer.reg_a = result
        computer.run_code()
        print(computer.get_output())
        if computer.get_output() == expected_output:
            return result

        while computer.get_output() != expected_output:
            result -= 1
            computer.reset()
            computer.reg_a = result
            computer.run_code()

        return result

        for i in range(len(first)):
            result = sum(((x * (8**i) for i, x in enumerate(first[: i + 1]))))
            print("Testing ", first[: i + 1], "=", result)
            computer.reset()
            computer.reg_a = result
            computer.run_code()
            print(expected_output)
            print(computer.get_output())

        if computer.get_output() == expected_output:
            return result
        else:
            return ""


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