import re

from typing import Dict
from typing import Iterable
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Tuple


NOP = 0b000
LDA = 0b001
STA = 0b010
ADD = 0b011
SUB = 0b100
BRZ = 0b101
JMP = 0b110
STP = 0b111

OPCODES = {
    : NOP,
    : LDA,
    : STA,
    : ADD,
    : SUB,
    : BRZ,
    : JMP,
    : STP,
}

RE_INSTRUCTION = re.compile(
    r
    r
    r
    rf
    r
    r
    r
    r
)


class AssemblySyntaxError(Exception):
    pass


class Instruction(NamedTuple):
    label: Optional[str]
    opcode: Optional[str]
    argument: Optional[str]
    comment: Optional[str]


def parse(assembly: str) -> Tuple[List[Instruction], Dict[str, int]]:
    instructions: List[Instruction] = []
    labels: Dict[str, int] = {}
    linenum: int = 0

    for line in assembly.split():
        match = RE_INSTRUCTION.match(line)

        if not match:
            raise AssemblySyntaxError(f)

        instructions.append(Instruction(**match.groupdict()))
        label = match.group(1)
        if label:
            labels[label] = linenum

        linenum += 1

    return instructions, labels


def compile(instructions: List[Instruction], labels: Dict[str, int]) -> Iterable[int]:
    for instruction in instructions:
        if instruction.argument is None:
            argument = 0
        elif instruction.argument.isnumeric():
            argument = int(instruction.argument)
        else:
            argument = labels[instruction.argument]

        if instruction.opcode:
            yield OPCODES[instruction.opcode] << 5 | argument
        else:
            yield argument


def run(bytecode: bytes) -> int:
    accumulator = 0
    program_counter = 0
    memory = list(bytecode)[:32] + [0 for _ in range(32 - len(bytecode))]

    while program_counter < 32:
        operation = memory[program_counter] >> 5
        argument = memory[program_counter] & 0b11111
        program_counter += 1

        if operation == NOP:
            continue
        elif operation == LDA:
            accumulator = memory[argument]
        elif operation == STA:
            memory[argument] = accumulator
        elif operation == ADD:
            accumulator = (accumulator + memory[argument])% 256
        elif operation == SUB:
            accumulator = (accumulator - memory[argument])% 256
        elif operation == BRZ:
            if accumulator == 0:
                program_counter = argument
        elif operation == JMP:
            program_counter = argument
        elif operation == STP:
            break
        else:
            raise Exception(f)

    return accumulator


SAMPLES = [
    ,
    ,
    ,
    ,
    ,
    ,
    ,
    ,
]


def main() -> None:
    for sample in SAMPLES:
        instructions, labels = parse(sample)
        bytecode = bytes(compile(instructions, labels))
        result = run(bytecode)
        print(result)


if __name__ == :
    main()