import sys

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
        }]
    }
}

labels = {}
lastLabel = ""

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):
        return False
    if (arg[0] != "R"):
        return False
    try:
        val = int(arg[1])
        if (0 <= val <= 3):
            return True
    except:
        pass
    return False

def testArgIsValue(arg):
    # Test 0 - 255
    try:
        val = int(arg)
        if (0 <= val <= 255):
            return True
    except:
        pass
    
    # Test 'a' 'A' '0'
    if (len(arg) == 3):
        if (arg[0] == arg[2] == "'"):
            if ((ord('a') <= ord(arg[1]) <= ord('z')) or (ord('A') <= ord(arg[1]) <= ord('Z')) or (ord('0') <= ord(arg[1]) <= ord('9'))):
                return True

    return False

def testArgIsLabel(arg, twoDotsIncluded = False):
    if (len(arg) == 0):
        return False
    if (arg[0] != "_"):
        return False
    if (twoDotsIncluded):
        if (arg[-1] != ":"):
            return False
        if (set(arg[1:-1]) <= set("abcdefghijklmnopqrstuvwxyz0123456789")):
            return True
    else:
        if (set(arg[1:]) <= set("abcdefghijklmnopqrstuvwxyz0123456789")):
            return True
    return False



def convertInsDB(args):
    value = valueToInt(args[0])
    return {"opcode": [value], "DB": True}

def convertInsCALL(args):

    return {"opcode": [0b00000000, "label"], "label": args[0], "offset": 0}

def convertInsRET(args):

    return {"opcode": [0b10000000]}

def convertInsJMP(args):

    return {"opcode": [0b01000000, "label"], "label": args[0], "offset": 0}

def convertInsJLT(args):

    return {"opcode": [0b11000000, "label"], "label": args[0], "offset": 0}

def convertInsJEQ(args):

    return {"opcode": [0b00100000, "label"], "label": args[0], "offset": 0}

def convertInsPUSH(args):
    idReg0 = registerToDec(args[0])
    return {"opcode": [0b10100000 | idReg0]}

def convertInsPOP(args):
    idReg0 = registerToDec(args[0])
    return {"opcode": [0b01100000 | idReg0]}

def convertInsMOV(args):
    idReg0 = registerToDec(args[0])
    print("idReg0", idReg0)
    if (testArgIsRegister(args[1])):
        idReg1 = registerToDec(args[1])
        print("idReg0", idReg1)
        return {"opcode": [0b01010000 | (idReg0 << 2) | (idReg1)]}
    value = valueToInt(args[1])
    return {"opcode": [0b11100000 | (idReg0), value]}


def convertInsSUB(args):
    idReg0 = registerToDec(args[0])
    print("idReg0", idReg0)
    if (testArgIsRegister(args[1])):
        idReg1 = registerToDec(args[1])
        print("idReg0", idReg1)
        return {"opcode": [0b11010000 | (idReg0 << 2) | (idReg1)]}
    value = valueToInt(args[1])
    return {"opcode": [0b00010000 | (idReg0), value]}

def convertInsCMP(args):
    idReg0 = registerToDec(args[0])
    print("idReg0", idReg0)
    if (testArgIsRegister(args[1])):
        idReg1 = registerToDec(args[1])
        print("idReg0", idReg1)
        return {"opcode": [0b00110000 | (idReg0 << 2) | (idReg1)]}
    value = valueToInt(args[1])
    return {"opcode": [0b10010000 | (idReg0), value]}

def convertInsLDR(args):
    idReg0 = registerToDec(args[0])
    idReg1 = registerToDec(args[1])
    return {"opcode": [0b10110000 | (idReg0 << 2) | (idReg1), "label"], "label": args[2], "offset": 0, "DB_Update": True}

def convertInsSTR(args):
    idReg0 = registerToDec(args[0])
    idReg1 = registerToDec(args[1])
    return {"opcode": [0b01110000 | (idReg0 << 2) | (idReg1), "label"], "label": args[2], "offset": 0, "DB_Update": True}

def convertInsOUT(args):
    idReg0 = registerToDec(args[0])
    return {"opcode": [0b11110000 | idReg0]}

def convertInsTIM(args):
    value = valueToInt(args[0])
    return {"opcode": [0b11111000, value]}

def testArg(arg, insArg):
    valid = False
    # Test for isRegister
    if (insArg["isRegister"] and testArgIsRegister(arg)):
        valid = True

    # Test for isValue
    if (insArg["isValue"] and testArgIsValue(arg)):
        valid = True

    # Test for isLabel
    if (insArg["isLabel"] and testArgIsLabel(arg)):
        valid = True

    if (not valid):
        print(f"ERROR : Arg {arg} not valid !")
        exit(1)


    pass

def decodeInstruction(args, ins):
    for i in range(0, len(args)):
        testArg(args[i], ins["args"][i])

    if (ins["ins"] == "DB"):
        return convertInsDB(args)
    elif (ins["ins"] == "CALL") :
        return convertInsCALL(args)
    elif (ins["ins"] == "RET") :
        return convertInsRET(args)
    elif (ins["ins"] == "JMP") :
        return convertInsJMP(args)
    elif (ins["ins"] == "JLT") :
        return convertInsJLT(args)
    elif (ins["ins"] == "JEQ") :
        return convertInsJEQ(args)
    elif (ins["ins"] == "PUSH") :
        return convertInsPUSH(args)
    elif (ins["ins"] == "POP") :
        return convertInsPOP(args)
    elif (ins["ins"] == "MOV") :
        return convertInsMOV(args)
    elif (ins["ins"] == "SUB") :
        return convertInsSUB(args)
    elif (ins["ins"] == "CMP") :
        return convertInsCMP(args)
    elif (ins["ins"] == "LDR") :
        return convertInsLDR(args)
    elif (ins["ins"] == "STR") :
        return convertInsSTR(args)
    elif (ins["ins"] == "OUT") :
        return convertInsOUT(args)
    elif (ins["ins"] == "TIM") :
        return convertInsTIM(args)


    pass

def decodeLine(line, PC):
    global lastLabel, labels
    commentPos = line.find(";")
    if (commentPos != -1):
        line = line[:line.find(";")]
    line = line.strip()
    #print(">" + line + "<")

    args = line.split(" ")
    args = [i for i in args if i]
    if (len(args) == 0):
        return
    INS = args[0]
    args = args[1:]
    #print(args)
    
    if (testArgIsLabel(INS, twoDotsIncluded=True)):
        lastLabel = INS[:-1]
        labels[lastLabel] = PC
        return

    instruction = None
    try:
        instruction = instructions[INS]
    except:
        print("ERROR : Bad instruction :", INS)
        exit(1)
    #print(instruction)
    if (len(args) != len(instruction["args"])):
        print(f"ERROR : Bad argument count. Excpected {len(instruction['args'])}, got {len(args)}")
        exit(1)
    
    return decodeInstruction(args, instruction)


def assemble(path):
    global lastLabel, labels
    PC = 0
    assemble1st = []
    bytecode = []
    with open(path, "r") as file:
        # 1er lecture, pre-compilation
        for line in file:
            print(line, end="")
            ret = decodeLine(line, PC)
            if (ret != None):
                PC += len(ret["opcode"])
                assemble1st.append(ret)
            print(" ==> ", ret)
        print("\n\n\n\n\n\n")
        print(assemble1st)

        print("Labels : ", labels)

        # Expansion des labels
        for item in assemble1st:
            if ("label" in item):
                labelIndex = labels[item["label"]]
                for index in range(len(item["opcode"])):
                    if (item["opcode"][index] == "label"):
                        item["opcode"][index] = labelIndex
            bytecode.extend(item["opcode"])

        print("\n\n\n\n\n\n")
        print(assemble1st)
        print(bytecode)
    return bytecode


if (__name__ == "__main__"):
    path = ""
    args = sys.argv
    if (len(args) > 1):
        path = args[1]
    else:
        print("NEED PATH !!!")
        exit(0)
    print(path)


    code = assemble(path)
    with open("out.bin", "wb") as file:
        file.write(bytes(code))
    exit(0)

