from sys import stdin, stdout
from os import PathLike
from interpreter import Interpreter, Registers, opcode
from pathlib import Path


BrainFuckByteCode = {
    ">": 1,
    "<": 2,
    "+": 3,
    "-": 4,
    ".": 5,
    ",": 6,
    "[": 7,
    "]": 8,
}


class BrainFuckByteCodeCompiler:
    def load_file(self, file: PathLike):
        self.source_file: PathLike = file
        self.load(bytearray(Path(file).read_bytes()))

    def load(self, data: bytearray):
        self.source_data: bytearray = data

    def compile(self) -> bytearray:
        return bytearray(
            BrainFuckByteCode[c]
            for byte in self.source_data
            if (c := chr(byte)) in BrainFuckByteCode
        )


class BrainFuckRegisters(Registers):
    def __init__(self) -> None:
        super().__init__()
        self.memory_pointer: int = 0


class BrainFuckInterpreter(Interpreter):
    def __init__(self, bits=8, memsize=None) -> None:
        super().__init__(bits, memsize)
        self._registers: BrainFuckRegisters = BrainFuckRegisters()

    def load_program(self, program: bytearray) -> None:
        super().load_program(program)
        self._registers.memory_pointer = len(program) + 1

    @opcode(BrainFuckByteCode[">"])
    def forward(self):
        self._registers.memory_pointer += 1

    @opcode(BrainFuckByteCode["<"])
    def backward(self):
        self._registers.memory_pointer -= 1

    @opcode(BrainFuckByteCode["+"])
    def increment(self):
        val = self._memory[self._registers.memory_pointer]
        self._memory[self._registers.memory_pointer] = (val + 1) % 256

    @opcode(BrainFuckByteCode["-"])
    def decrement(self):
        val = self._memory[self._registers.memory_pointer]
        self._memory[self._registers.memory_pointer] = (val - 1) & 255

    @opcode(BrainFuckByteCode["."])
    def io_out(self):
        stdout.write(chr(self._memory[self._registers.memory_pointer]))

    @opcode(BrainFuckByteCode[","])
    def io_in(self):
        stdout.flush()
        self._memory[self._registers.memory_pointer] = ord(stdin.read(1))

    @opcode(BrainFuckByteCode["["])
    def loop_start(self):
        # Ignore opcode if mem_pointer == 0
        if self._memory[self._registers.memory_pointer] > 0:
            pass
        elif self._memory[self._registers.memory_pointer] == 0:
            self._registers.program_counter = self.findNext() - 1

    @opcode(BrainFuckByteCode["]"])
    def start_end(self):
        self._registers.program_counter = self.findNext(-1)

    def findNext(self, incrementValue: int = 1) -> int:
        search_pointer = self._registers.program_counter + incrementValue
        bracket_counter = 1
        while bracket_counter != 0:
            opcode = self._memory[search_pointer]
            if opcode == BrainFuckByteCode["["]:
                bracket_counter += incrementValue
            elif opcode == BrainFuckByteCode["]"]:
                bracket_counter -= incrementValue
            search_pointer += incrementValue
        return search_pointer


def bf(program_path: PathLike, bits=8, mem=None) -> BrainFuckInterpreter:
    bf_c = BrainFuckByteCodeCompiler()
    bf_c.load_file(program_path)
    program = bf_c.compile()

    bf_e = BrainFuckInterpreter(bits=bits, memsize=mem)
    bf_e.load_program(program)
    bf_e.run()

    return bf_e


if __name__ == "__main__":
    bf(Path("examples", "HelloWorld.bf"))
    bf(Path("examples", "GameOfLife.bf"), bits=16, mem=2**16)
    bf(Path("examples", "Mandlebrot.bf"), bits=16, mem=2**16)