import memory

from memory import Convention
from memory import DataType
from memory import Function


class AsmFunction(Function):
    def __init__(self, op_codes, convention=Convention.CDECL, args=[], return_type=DataType.VOID):
        if not op_codes:
            raise ValueError('At least one OP code is required.')

        self.op_codes = op_codes

        # Allocate memory
        super().__init__(
            memory.alloc(len(op_codes), False).address,
            convention,
            args,
            return_type
        )

        # Make the buffer executable
        self.unprotect(len(op_codes))

        # Write OP codes
        for offset, op_code in enumerate(op_codes):
            self.set_uchar(op_code, offset)


multiply = AsmFunction(
    [
        0x55,                   # push    ebp
        0x8B, 0xEC,             # mov     ebp, esp
        0x8B, 0x45, 0x08,       # mov     eax, [ebp+8]
        0x0F, 0xAF, 0x45, 0x0C, # imul    eax, [ebp+0Ch]
        0x5D,                   # pop     ebp
        0xC3                    # retn
    ],
    args=[DataType.INT, DataType.INT],
    return_type=DataType.INT
)

# Output: 21
print(multiply(7, 3))

multiply.dealloc()