import sys

# --- Configuration des instructions (Inchangée) ---
instructions = {
    "DB": {"ins": "DB", "args": [{"isRegister": False, "isValue": True, "isLabel": False}]},
    "CALL": {"ins": "CALL", "args": [{"isRegister": False, "isValue": False, "isLabel": True}]},
    "RET": {"ins": "RET", "args": []},
    "JMP": {"ins": "JMP", "args": [{"isRegister": False, "isValue": False, "isLabel": True}]},
    "JLT": {"ins": "JLT", "args": [{"isRegister": False, "isValue": False, "isLabel": True}]},
    "JEQ": {"ins": "JEQ", "args": [{"isRegister": False, "isValue": False, "isLabel": True}]},
    "PUSH": {"ins": "PUSH", "args": [{"isRegister": True, "isValue": False, "isLabel": False}]},
    "POP": {"ins": "POP", "args": [{"isRegister": True, "isValue": False, "isLabel": False}]},
    "MOV": {"ins": "MOV", "args": [{"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": True, "isValue": True, "isLabel": False}]},
    "SUB": {"ins": "SUB", "args": [{"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": True, "isValue": True, "isLabel": False}]},
    "CMP": {"ins": "CMP", "args": [{"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": True, "isValue": True, "isLabel": False}]},
    "LDR": {"ins": "LDR", "args": [{"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": False, "isValue": False, "isLabel": True}]},
    "STR": {"ins": "STR", "args": [{"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": True, "isValue": False, "isLabel": False}, {"isRegister": False, "isValue": False, "isLabel": True}]},
    "OUT": {"ins": "OUT", "args": [{"isRegister": True, "isValue": False, "isLabel": False}]},
    "TIM": {"ins": "TIM", "args": [{"isRegister": False, "isValue": True, "isLabel": False}]}
}

# --- Fonctions Utilitaires (Inchangées) ---
def valueToInt(arg):
    try: return int(arg)
    except: return ord(arg[1])

def registerToDec(reg): return int(reg[1])

def testArgIsRegister(arg):
    if len(arg) != 2 or arg[0] != "R": return False
    try: return 0 <= int(arg[1]) <= 3
    except: return False

def testArgIsValue(arg):
    try:
        if 0 <= int(arg) <= 255: return True
    except: pass
    if len(arg) == 3 and arg[0] == arg[2] == "'": return True
    return False

def testArgIsLabel(arg, twoDotsIncluded=False):
    if not arg or arg[0] != "_": return False
    body = arg[1:-1] if twoDotsIncluded else arg[1:]
    if twoDotsIncluded and arg[-1] != ":": return False
    return all(c in "abcdefghijklmnopqrstuvwxyz0123456789_" for c in body)

# --- Fonctions de Conversion (Inchangées) ---
def convertInsDB(args): return {"opcode": [valueToInt(args[0])], "is_db": True}
def convertInsCALL(args): return {"opcode": [0x00, "label"], "label": args[0]}
def convertInsRET(args): return {"opcode": [0x80]}
def convertInsJMP(args): return {"opcode": [0x40, "label"], "label": args[0]}
def convertInsJLT(args): return {"opcode": [0xC0, "label"], "label": args[0]}
def convertInsJEQ(args): return {"opcode": [0x20, "label"], "label": args[0]}
def convertInsPUSH(args): return {"opcode": [0xA0 | registerToDec(args[0])]}
def convertInsPOP(args): return {"opcode": [0x60 | registerToDec(args[0])]}
def convertInsMOV(args):
    idReg0 = registerToDec(args[0])
    if testArgIsRegister(args[1]): return {"opcode": [0x50 | (idReg0 << 2) | registerToDec(args[1])]}
    return {"opcode": [0xE0 | idReg0, valueToInt(args[1])]}
def convertInsSUB(args):
    idReg0 = registerToDec(args[0])
    if testArgIsRegister(args[1]): return {"opcode": [0xD0 | (idReg0 << 2) | registerToDec(args[1])]}
    return {"opcode": [0x10 | idReg0, valueToInt(args[1])]}
def convertInsCMP(args):
    idReg0 = registerToDec(args[0])
    if testArgIsRegister(args[1]): return {"opcode": [0x30 | (idReg0 << 2) | registerToDec(args[1])]}
    return {"opcode": [0x90 | idReg0, valueToInt(args[1])]}
def convertInsLDR(args):
    return {"opcode": [0xB0 | (registerToDec(args[0]) << 2) | registerToDec(args[1]), "label"], "label": args[2]}
def convertInsSTR(args):
    return {"opcode": [0x70 | (registerToDec(args[0]) << 2) | registerToDec(args[1]), "label"], "label": args[2]}
def convertInsOUT(args): return {"opcode": [0xF0 | registerToDec(args[0])]}
def convertInsTIM(args): return {"opcode": [0xF8, valueToInt(args[0])]}

# --- Assembleur ---
def assemble(path):
    labels = {}
    code_elements = []
    data_elements = []
    
    with open(path, "r") as f:
        lines = f.readlines()

    current_pending_labels = []

    for line_num, line in enumerate(lines, 1):
        line = line.split(";")[0].strip()
        if not line: continue
        
        parts = line.split()
        
        # Traitement des labels (on peut en avoir plusieurs de suite)
        while parts and testArgIsLabel(parts[0], True):
            label_name = parts[0][:-1]
            current_pending_labels.append(label_name)
            parts = parts[1:] # On retire le label et on continue sur la même ligne
            
        if not parts: continue # Ligne ne contenait que des labels
        
        instr_name = parts[0]
        args = parts[1:]
        
        try:
            res = globals()[f"convertIns{instr_name}"](args)
            res["attached_labels"] = current_pending_labels
            current_pending_labels = []
            
            if res.get("is_db"):
                data_elements.append(res)
            else:
                code_elements.append(res)
        except Exception as e:
            print(f"ERROR Line {line_num}: {instr_name} -> {e}")
            sys.exit(1)

    # Si des labels traînent à la toute fin du fichier
    if current_pending_labels:
        # On les attache à un élément factice à la fin de la data
        data_elements.append({"opcode": [], "attached_labels": current_pending_labels, "is_db": True})

    # CALCUL DES ADRESSES FINALES
    final_labels = {}
    current_pc = 0
    
    # 1. On passe sur le CODE d'abord
    for item in code_elements:
        for lbl in item["attached_labels"]:
            final_labels[lbl] = current_pc
        current_pc += len(item["opcode"])
        
    # 2. On passe sur la DATA ensuite
    for item in data_elements:
        for lbl in item["attached_labels"]:
            final_labels[lbl] = current_pc
        current_pc += len(item["opcode"])

    # GÉNÉRATION DU BYTECODE
    final_bytecode = []
    # Ordre : Code puis Data
    for item in code_elements + data_elements:
        for op in item["opcode"]:
            if op == "label":
                label_target = item["label"]
                if label_target not in final_labels:
                    print(f"ERROR: Label '{label_target}' missing!")
                    sys.exit(1)
                final_bytecode.append(final_labels[label_target])
            else:
                final_bytecode.append(op)
                
    return final_bytecode, final_labels

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python asm.py <file.asm>")
        sys.exit(1)
        
    path = sys.argv[1]
    bytecode, labels_map = assemble(path)

    # Affichage propre
    print("\n" + "="*50)
    print(f" ASSEMBLY PREVIEW: {path}")
    print("="*50)
    print(f"{'ADDR':<7} | {'HEX':<5} | {'BINARY':<10}")
    print("-" * 30)
    for i, b in enumerate(bytecode):
        # On cherche si un label pointe ici pour l'afficher
        lbl_str = ""
        for name, addr in labels_map.items():
            if addr == i: lbl_str += f" ({name})"
        
        print(f"0x{i:02X}    | {b:02x}   | {b:08b} {lbl_str}")
    print("="*50 + "\n")

    with open(path + ".bin", "wb") as file:
        file.write(bytes(bytecode))
    print(f"Success: {path}.bin generated.")
