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


class Command(str, Enum):
    COPY = "cpy"
    INCREASE = "inc"
    DECREASE = "dec"
    JUMP_NON_ZERO = "jnz"


class Computer:
    def __init__(
        self, code: list[tuple[Command, str | int, ...]], reg_a: int = 0, reg_b: int = 0, reg_c: int = 0, reg_d: int = 0
    ):
        self.reg = {
            "a": reg_a,
            "b": reg_b,
            "c": reg_c,
            "d": reg_d,
        }
        self.code = code
        self.instr_ptr = 0

    def run(self):
        while self.instr_ptr < len(self.code):
            cmd, *args = self.code[self.instr_ptr]
            if cmd == Command.COPY:
                if isinstance(args[0], int):
                    self.reg[args[1]] = args[0]
                else:
                    self.reg[args[1]] = self.reg[args[0]]
                self.instr_ptr += 1
            elif cmd == Command.INCREASE:
                self.reg[args[0]] += 1
                self.instr_ptr += 1
            elif cmd == Command.DECREASE:
                self.reg[args[0]] -= 1
                self.instr_ptr += 1
            elif cmd == Command.JUMP_NON_ZERO:
                if isinstance(args[0], int):
                    check = args[0]
                else:
                    check = self.reg[args[0]]
                if check:
                    self.instr_ptr += args[1]
                else:
                    self.instr_ptr += 1


class Day(AOCDay):
    inputs = [
        [
            (42, "input12_test"),
            (318083, "input12"),
        ],
        [
            (9227737, "input12"),
        ],
    ]

    def parse_input(self) -> list[tuple[Command, str | int, ...]]:
        code = []
        for line in self.getInput():
            args = line.split()
            cmd = Command(args[0])
            args = [int(arg) if arg not in "abcd" else arg for arg in args[1:]]
            code.append((cmd, *args))

        return code

    def part1(self) -> Any:
        computer = Computer(code=self.parse_input())
        computer.run()
        return computer.reg["a"]

    def part2(self) -> Any:
        computer = Computer(code=self.parse_input(), reg_c=1)
        computer.run()
        return computer.reg["a"]


if __name__ == "__main__":
    day = Day(2016, 12)
    day.run(verbose=True)