import re
import struct
from typing import List, Dict, Union


class PDP11Assembler:
    def __init__(self):
        self.labels: Dict[str, int] = {}
        self.current_address = 0o1000
        self.output_data: List[int] = []

        # Таблица команд с правильными опкодами
        self.commands = {
            'halt': {'opcode': 0o000000, 'args': []},
            'mov': {'opcode': 0o010000, 'args': ['ss', 'dd']},
            'add': {'opcode': 0o060000, 'args': ['ss', 'dd']},
        }

    def parse_operand(self, operand: str) -> Dict[str, Union[int, str]]:
        """Разбор операнда с правильными режимами адресации"""
        if not operand:
            return {'mode': 0, 'reg': 0, 'value': 0}

        # Регистровая адресация (R0-R7)
        if re.match(r'^[rR][0-7]$', operand):
            return {'mode': 0, 'reg': int(operand[1]), 'value': 0}

        # Немедленная адресация (#n)
        if operand.startswith('#'):
            value = int(operand[1:])
            return {'mode': 0o27, 'reg': 0, 'value': value}  # Режим 27, регистр 0

        return {'mode': 0, 'reg': 0, 'value': 0}

    def assemble_instruction(self, command: List[str]) -> List[Union[int, str]]:
        """Генерация машинного кода с правильным форматом"""
        if not command:
            return []

        cmd_name = command[0].lower()
        if cmd_name not in self.commands:
            return []

        cmd_info = self.commands[cmd_name]
        args = command[1:]

        if len(args) != len(cmd_info['args']):
            return []

        words = []
        opcode = cmd_info['opcode']

        if cmd_name == 'halt':
            words.append(0o000000)
            return words

        if cmd_name == 'mov':
            src = self.parse_operand(args[0])
            dst = self.parse_operand(args[1])

            # Правильное кодирование для MOV
            instruction = opcode | (src['mode'] << 6) | (dst['reg'] << 3) | dst['mode']
            words.append(instruction)

            if src['mode'] == 0o27:  # Немедленная адресация
                words.append(src['value'])

            return words

        if cmd_name == 'add':
            src = self.parse_operand(args[0])
            dst = self.parse_operand(args[1])

            # Правильное кодирование для ADD
            instruction = opcode | (src['reg'] << 6) | dst['reg']
            words.append(instruction)

            return words

        return []

    def first_pass(self, lines: List[str]):
        """Первый проход: сбор меток"""
        self.current_address = 0o1000
        for line in lines:
            line = line.strip()
            if not line or line.startswith(';'):
                continue

            if ':' in line:
                label, _ = line.split(':', 1)
                self.labels[label.strip()] = self.current_address
                line = line.split(':', 1)[1].strip()
                if not line:
                    continue

            parts = re.split(r'[,\s]+', line)
            cmd = parts[0].lower()

            if cmd in self.commands:
                self.current_address += 2
                if cmd == 'mov' and parts[1].startswith('#'):
                    self.current_address += 2

    def second_pass(self, lines: List[str]):
        """Второй проход: генерация кода"""
        self.current_address = 0o1000
        for line in lines:
            line = line.strip()
            if not line or line.startswith(';'):
                continue

            if ':' in line:
                line = line.split(':', 1)[1].strip()
                if not line:
                    continue

            parts = re.split(r'[,\s]+', line)
            machine_code = self.assemble_instruction(parts)

            for word in machine_code:
                self.output_data.append(word)
                self.current_address += 2

    def write_output_file(self, filename: str):
        """Запись в файл в нужном формате"""
        data_bytes = bytearray()
        for word in self.output_data:
            data_bytes.extend(struct.pack('<H', word))  # Little-endian

        with open(filename, 'w') as f:
            # Первая строка: адрес и размер
            f.write(f"0200 {len(data_bytes):04x}\n")

            # Байты машинного кода
            for byte in data_bytes:
                f.write(f"{byte:02x}\n")

    def assemble(self, input_filename: str, output_filename: str):
        """Основной метод ассемблирования"""
        with open(input_filename, 'r') as f:
            lines = f.readlines()

        self.first_pass(lines)
        self.second_pass(lines)
        self.write_output_file(output_filename)


# Пример использования
if __name__ == "__main__":
    assembler = PDP11Assembler()

    test_code = """
        mov #2, r0
        mov #3, r1
        add r0, r1
        halt
    """

    with open("test.asm", "w") as f:
        f.write(test_code)

    assembler.assemble("test.asm", "output.txt")
    print("Ассемблирование завершено. Результат в output.txt")