#!/usr/bin/env python3

from dataclasses import dataclass, field
from typing import Any, Callable, Optional

from aoc import Puzzle, run


@dataclass
class VM:
    ops: list[tuple[str, int]]
    ip: int = field(default=0)  # instruction pointer
    stopped: bool = field(default=False)

    acc: int = field(default=0)  # accumulator

    callback: Optional[Callable[["VM"], Any]] = field(default=None)

    @classmethod
    def from_string(cls, data: str) -> "VM":
        ops: list[tuple[str, int]] = []

        for line in data.splitlines():
            line = line.strip()
            op, arg = line.split()
            ops.append((op, int(arg)))

        return cls(ops=ops)

    def clone(self) -> "VM":
        return VM(ops=self.ops.copy())

    def reset(self):
        self.ip = 0
        self.stopped = False
        self.acc = 0
        self.callback = None

    def step(self):
        """Run one instruction."""
        if self.callback:
            self.callback(self)
        if self.stopped:
            return

        op, arg = self.ops[self.ip]
        func = getattr(self, f"op_{op}")
        keep_ip = func(arg)
        if not keep_ip:
            self.ip += 1

    def run(self):
        """Run the program until stopped."""
        while 0 <= self.ip < len(self.ops) and not self.stopped:
            self.step()

    def op_nop(self, arg):
        pass

    def op_acc(self, arg: int):
        self.acc += arg

    def op_jmp(self, arg: int):
        self.ip += arg
        return True


class Day08(Puzzle):
    test_data = [
        """nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6"""
    ]
    test_result_part1 = [5]
    test_result_part2 = [8]

    def prepare_data(self, data: str) -> VM:
        return VM.from_string(data)

    def run_part1(self, data: VM):
        visited_ips = set()
        result = None

        def check_visited_ips(vm):
            nonlocal result
            if vm.ip in visited_ips:
                vm.stopped = True
                result = vm.acc
            else:
                visited_ips.add(vm.ip)

        vm = data.clone()
        vm.callback = check_visited_ips

        vm.run()
        return result

    def run_part2(self, data):
        visited_ips = set()
        prev_fixes = set()
        fixed = False

        def fix_corruption(vm):
            nonlocal fixed
            ip = vm.ip
            looping = False

            if ip in visited_ips:
                looping = True
                if fixed:
                    # Try again.
                    vm.stopped = True

            visited_ips.add(ip)

            op, arg = vm.ops[ip]
            if op == "acc":
                return

            if looping:
                # What would happen if we changed the current op?
                new_ip, alt_ip, alt_op = -1, -1, ""
                if op == "nop":
                    new_ip = ip + 1
                    alt_ip = ip + arg
                    alt_op = "jmp"
                else:
                    new_ip = ip + arg
                    alt_ip = ip + 1
                    alt_op = "nop"

                # Would flipping this op get us out of the loop?
                if (
                    new_ip in visited_ips
                    and alt_ip not in visited_ips
                    and ip not in prev_fixes
                ):
                    # Bingo!
                    # print(f"{ip}: {op} {arg} --> {alt_op} {arg}")
                    vm.ops[ip] = (alt_op, arg)
                    fixed = True
                    prev_fixes.add(ip)

        vm = data.clone()
        vm.stopped = True
        while vm.stopped:
            visited_ips.clear()
            fixed = False

            vm = data.clone()
            vm.callback = fix_corruption
            vm.run()

        # Not stopped: that fix worked! Now let's run it uncorrupted.
        vm.reset()
        vm.run()

        return vm.acc

    def run_part2_bruteforce(self, data):
        visited_ips = set()

        def detect_loop(vm):
            if vm.ip in visited_ips:
                vm.stopped = True
            else:
                visited_ips.add(vm.ip)

        for pos in range(len(data.ops)):
            op, arg = data.ops[pos]
            if op not in ("nop", "jmp"):
                continue

            vm = data.clone()
            visited_ips.clear()
            new_op = "jmp" if op == "nop" else "nop"
            vm.ops[pos] = (new_op, arg)
            vm.callback = detect_loop
            vm.run()

            if not vm.stopped:
                # Program ran without being interrupted!
                return vm.acc


if __name__ == "__main__":
    run(obj=Day08())