#!/bin/env python3

import sys
import struct
import re

# ==========================================
# ISA DEFINITIONS
# ==========================================

REGISTERS = {
    'R0': 0, 'R1': 1, 'R2': 2, 'R3': 3,
    'R4': 4, 'R5': 5, 'SP': 6, 'PC': 7
}

OPCODES = {
    'NOP': 0x00, 'HALT': 0xFF,
    'SET': 0x01, 'MOV': 0x02,
    'LODB': 0x03, 'STOB': 0x04,
    'LODW': 0x05, 'STOW': 0x06,
    'ADD': 0x10, 'SUB': 0x11,
    'MUL': 0x17, 'DIV': 0x18,
    'INC': 0x19,
    'AND': 0x12, 'OR':  0x13,
    'XOR': 0x14, 'NOT': 0x15,
    'CMP': 0x16,
    'JMP': 0x20, 'JZ':  0x21,
    'JNZ': 0x22, 'JG':  0x23,
    'JL':  0x24, 'JLE': 0x25,
    'JGE': 0x26,
    'PUSH': 0x30, 'POP': 0x31,
    'CALL': 0x32, 'RET': 0x33,
    'HW_SCAN':   0xF0,
    'HW_SELECT': 0xF1,
    'HW_TYPE':   0xF2,
    'STI': 0xFB, 'CLI': 0xFC, 'IRET': 0xFD
}

MULTI_WORD_INSTRUCTIONS = [
    0x01, 
    0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26,
    0x32
]

# ==========================================
# ASSEMBLER LOGIC
# ==========================================

def is_register(token):
    return token.upper() in REGISTERS

def parse_register(token):
    return REGISTERS[token.upper()]

def parse_number(token):
    try:
        return int(token, 0)
    except ValueError:
        return None

def pack_instruction(opcode, dest, src):
    low_byte = (dest << 4) | (src & 0x0F)
    combined = (opcode << 8) | low_byte
    return struct.pack('<H', combined)

def pack_immediate(val):
    return struct.pack('<H', val & 0xFFFF)

def main(input_file, output_file):
    print(f"Assembling {input_file} -> {output_file}...")

    with open(input_file, 'r') as f:
        lines = f.readlines()

    clean_lines = []
    symbol_table = {}
    constants = {}
    
    # Start at 0x2000 if compiling OS for User Space, or 0x0000 if Kernel.
    # For the KERNEL you just wrote, set this to 0x0000.
    # For USER programs (like 'game'), set this to 0x2000.
    current_address = 0x0000 

    # --- PASS 1: SYMBOLS & STRING SIZING ---
    for line_idx, line in enumerate(lines):
        line = line.split(';')[0].strip()
        if not line: continue

        # 1. HANDLE STRINGS (.STRING "Text")
        # We process this BEFORE tokenizing to preserve spaces in quotes.
        if '.STRING' in line.upper():
            # Check for label "MSG: .STRING"
            if ':' in line:
                parts = line.split(':', 1)
                label_name = parts[0].strip()
                symbol_table[label_name] = current_address
                line = parts[1].strip() # Remain with ".STRING ..."

            # Extract content between quotes
            start_quote = line.find('"')
            end_quote = line.rfind('"')
            
            if start_quote != -1 and end_quote != -1 and end_quote > start_quote:
                content = line[start_quote+1 : end_quote]
                # Length = characters + 1 null terminator
                size = len(content) + 1
                
                # Store special tuple for Pass 2
                clean_lines.append((line_idx, ['ALL_STRING', content]))
                current_address += size
                continue
            else:
                print(f"Error Line {line_idx+1}: Malformed string.")
                sys.exit(1)

        # 2. CONSTANTS
        if '=' in line:
            parts = line.split('=')
            name = parts[0].strip().upper()
            val = parse_number(parts[1].strip())
            constants[name] = val
            continue

        # 3. LABELS
        if line.endswith(':'):
            label_name = line[:-1].strip()
            symbol_table[label_name] = current_address
            continue
        
        # 4. NORMAL INSTRUCTION
        tokens = [t for t in re.split(r'[ ,\[\]]+', line) if t]
        mnemonic = tokens[0].upper()
        
        if mnemonic not in OPCODES:
            print(f"Error Line {line_idx+1}: Unknown Opcode '{mnemonic}'")
            sys.exit(1)

        opcode = OPCODES[mnemonic]
        current_address += 2 
        if opcode in MULTI_WORD_INSTRUCTIONS:
            current_address += 2

        clean_lines.append((line_idx, tokens))

    # --- PASS 2: CODE GENERATION ---
    binary_out = bytearray()

    for line_idx, tokens in clean_lines:
        mnemonic = tokens[0].upper()
        
        # --- HANDLE STRING GENERATION ---
        if mnemonic == 'ALL_STRING':
            content = tokens[1]
            # Convert chars to bytes
            for char in content:
                binary_out.append(ord(char))
            binary_out.append(0) # Null terminator
            continue

        # --- HANDLE INSTRUCTIONS ---
        opcode = OPCODES[mnemonic]
        dest = 0
        src = 0
        immediate = None

        args = tokens[1:]
        
        def resolve_value(token):
            val = parse_number(token)
            if val is not None: return val
            if token.upper() in constants: return constants[token.upper()]
            if token in symbol_table: return symbol_table[token]
            return None

        # Two Registers
        if len(args) == 2 and is_register(args[0]) and is_register(args[1]):
            reg_a = parse_register(args[0])
            reg_b = parse_register(args[1])
            if mnemonic in ['STOB', 'STOW']:
                dest = reg_b
                src = reg_a
            else:
                dest = reg_a
                src = reg_b

        # One Register, One Immediate
        elif len(args) == 2 and is_register(args[0]):
            dest = parse_register(args[0])
            immediate = resolve_value(args[1])
            if immediate is None:
                 print(f"Error Line {line_idx+1}: Unknown symbol '{args[1]}'")
                 sys.exit(1)

        # One Operand
        elif len(args) == 1:
            if is_register(args[0]):
                if mnemonic in ['POP', 'INC', 'NOT']:
                    dest = parse_register(args[0])
                elif mnemonic == 'HW_SELECT':
                    src = parse_register(args[0])
                else:
                    src = parse_register(args[0])
            else:
                immediate = resolve_value(args[0])
                if immediate is None:
                    print(f"Error Line {line_idx+1}: Unknown label '{args[0]}'")
                    sys.exit(1)

        binary_out += pack_instruction(opcode, dest, src)
        
        if opcode in MULTI_WORD_INSTRUCTIONS:
            binary_out += pack_immediate(immediate)

    with open(output_file, 'wb') as f:
        f.write(binary_out)
    
    print(f"Success! {len(binary_out)} bytes.")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python assembler.py <input.asm> [output.bin]")
    else:
        main(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else "program.bin")
