import inspect
from typing import Any, Callable, Final, Iterator, Mapping

from ..base import BaseSolver


class Instruction:
    def __init__(self, fn: Callable[..., None]):
        self._fn = fn

        args = inspect.getfullargspec(fn)

        self._argtypes = [args.annotations[arg] for arg in args.args[1:]]

    def __call__(self, args: tuple[str, ...]):
        self._fn(
            *(argtype(arg) for arg, argtype in zip(args, self._argtypes, strict=True))
        )


class Machine:
    def __init__(
        self, instructions: list[str], registers: dict[str, int] = {"a": 0, "b": 1}
    ):
        self.instructions: Final = [
            (part[0], tuple(arg.strip() for arg in " ".join(part[1:]).split(",")))
            for instruction in instructions
            if (part := instruction.split())
        ]

        self._fns = {
            name: Instruction(getattr(self, name))
            for name in ("hlf", "tpl", "inc", "jmp", "jie", "jio")
        }

        self._registers = registers.copy()
        self._ip = 0

    @property
    def registers(self) -> Mapping[str, int]:
        return self._registers

    @property
    def ip(self) -> int:
        return self._ip

    def reset(self, registers: dict[str, int] = {"a": 0, "b": 0}):
        self._registers = registers.copy()
        self._ip = 0

    def hlf(self, register: str):
        self._registers[register] //= 2
        self._ip += 1

    def tpl(self, register: str):
        self._registers[register] *= 3
        self._ip += 1

    def inc(self, register: str):
        self._registers[register] += 1
        self._ip += 1

    def jmp(self, offset: int):
        self._ip += offset
        assert 0 <= self._ip < len(self.instructions)

    def jie(self, register: str, offset: int):
        if self._registers[register] % 2 == 0:
            self._ip += offset
        else:
            self._ip += 1

    def jio(self, register: str, offset: int):
        if self._registers[register] == 1:
            self._ip += offset
        else:
            self._ip += 1

    def _exec(self) -> bool:
        # execute next instruction
        if self._ip >= len(self.instructions):
            return False

        ins, args = self.instructions[self._ip]

        if ins not in self._fns:
            return False

        self._fns[ins](args)
        return True

    def run(self):
        while self._exec():
            ...
        return self.registers


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        machine = Machine(input.splitlines())

        registers = machine.run()
        yield registers["b"]

        machine.reset({"a": 1, "b": 0})
        registers = machine.run()
        yield registers["b"]