from ScEpTIC.AST.builtins.builtin import Builtin
from ScEpTIC.AST.elements.value import Value

"""
LLVM library functions
"""

def get_unaligned_cells(cells, nbytes, stack_limit_address, stack_prefix):
    """
    ScEpTIC align memory to single bytes.
    To make llvm mem-based functions work, we need to count the number of cells instead of using the memory size.
    Further, when allocating memory without initializing the cells with a value, ScEpTIC creates a single cell of the entire data structure size.
    To ensure we do not smash the stack, this function limits the cells to EBP-2 cells.
    This is possible as these Builtins are function calls and allocate a new stack frame.
    :param cells: List of cells
    :param nbytes: number of requested bytes
    :param stack_limit_address: start address of the new frame (EBP - 2 cells, for saved PC and EBP)
    :param stack_prefix: address prefix of the stack
    """

    # LLVM uses cell's alignment to calculate nbytes
    if not isinstance(cells, list) or len(cells) == 0:
        return cells

    alignment = cells[0].alignment
    if alignment is not None and alignment > 0:
        # Calculate real number of cells
        n_cells = nbytes // alignment
        # Resize cells using ScEpTIC no-alignment policy
        cells = cells[:n_cells]

    if len(cells) > 0 and cells[0].address_prefix == stack_prefix:
        for idx, cell in enumerate(cells):
            if cell.address >= stack_limit_address:
                return cells[:idx]

    return cells

class msp430_dma_word_copy(Builtin):
    """
    msp430 __dma_word_copy() builtin
    """

    def get_val(self):
        args = self.args
        destination = args[0]
        source = args[1]
        # compiled for x86 -> word size 32 bit -> 4 bytes
        nbytes = args[2]

        # Stack limit address
        stack_limit_address = self._vmstate.register_file.ebp - 2 * (self._vmstate.memory.address_dimension // 8)
        stack_prefix = self._vmstate.memory.stack.address_prefix

        cells = self._vmstate.memory.get_cells_from_address(source, nbytes, False)
        # No cells alignment to be considered
        self._vmstate.memory.simulate_mmu_read_on_cells(cells)

        self._vmstate.memory.set_cells_from_address(destination, cells)

class memcpy(Builtin):
    """
    memcpy() builtin
    """

    def get_val(self):
        args = self.args
        destination = args[0]
        source = args[1]
        nbytes = args[2]

        # Stack limit address
        stack_limit_address = self._vmstate.register_file.ebp - 2 * (self._vmstate.memory.address_dimension // 8)
        stack_prefix = self._vmstate.memory.stack.address_prefix

        cells = self._vmstate.memory.get_cells_from_address(source, nbytes, False)
        cells = get_unaligned_cells(cells, nbytes, stack_limit_address, stack_prefix)
        self._vmstate.memory.simulate_mmu_read_on_cells(cells)

        self._vmstate.memory.set_cells_from_address(destination, cells)


class memmove(Builtin):
    """
    memmove() builtin
    """

    def get_val(self):
        args = self.args
        destination = args[0]
        source = args[1]
        nbytes = args[2]

        # Stack limit address
        stack_limit_address = self._vmstate.register_file.ebp - 2 * (self._vmstate.memory.address_dimension // 8)
        stack_prefix = self._vmstate.memory.stack.address_prefix

        cells = self._vmstate.memory.get_cells_from_address(source, nbytes, False)
        cells = get_unaligned_cells(cells, nbytes, stack_limit_address, stack_prefix)
        self._vmstate.memory.simulate_mmu_read_on_cells(cells)
        self._vmstate.memory.set_cells_from_address(destination, cells)

        cells = self._vmstate.memory.get_cells_from_address(destination, nbytes, False)
        cells = get_unaligned_cells(cells, nbytes, stack_limit_address, stack_prefix)
        self._vmstate.memory.simulate_mmu_read_on_cells(cells)
        for cell in cells:
            self._vmstate.memory.write(cell.absolute_address, cell.dimension, None, False)


class memset(Builtin):
    """
    memset() builtin
    """

    def get_val(self):
        args = self.args
        destination = args[0]
        nbytes = args[2]
        value = args[1]

        # Stack limit address
        stack_limit_address = self._vmstate.register_file.ebp - 2 * (self._vmstate.memory.address_dimension // 8)
        stack_prefix = self._vmstate.memory.stack.address_prefix

        cells = self._vmstate.memory.get_cells_from_address(destination, nbytes, False)
        # Keep this here to ensure data structures works
        # TODO: change alignment to the one of the operation, and not the one of the cell!!
        cells = get_unaligned_cells(cells, nbytes, stack_limit_address, stack_prefix)
        self._vmstate.memory.simulate_mmu_read_on_cells(cells)
        for cell in cells:
            self._vmstate.memory.write(cell.absolute_address, cell.dimension, value, False)

"""
Definitions for the various versions of llvm.memcpy, llvm.memmove, llvm.memset
"""
# Called on linking, as it depends on vmstate
def address_dimension_16():
    dim = f"i{Builtin._vmstate.memory.address_dimension}"
    return f"{dim}, {dim}, i16"

def address_dimension_32():
    dim = f"i{Builtin._vmstate.memory.address_dimension}"
    return f"{dim}, {dim}, i32"

def address_dimension_64():
    dim = f"i{Builtin._vmstate.memory.address_dimension}"
    return f"{dim}, {dim}, i64"

# in llvm 5, source destination len align isvolatile, in llvm7 align is missing.
# To support every version, it is defined only as if it has 3 arguments.
memcpy.define_builtin('llvm.memcpy.p0i8.p0i8.i16', address_dimension_16, 'void')
memcpy.define_builtin('llvm.memcpy.p0i8.p0i8.i32', address_dimension_32, 'void')
memcpy.define_builtin('llvm.memcpy.p0i8.p0i8.i64', address_dimension_64, 'void')

memmove.define_builtin('llvm.memmove.p0i8.p0i8.i16', address_dimension_16, 'void')
memmove.define_builtin('llvm.memmove.p0i8.p0i8.i32', address_dimension_32, 'void')
memmove.define_builtin('llvm.memmove.p0i8.p0i8.i64', address_dimension_64, 'void')

memset.define_builtin('llvm.memset.p0i8.i16', address_dimension_16, 'void')
memset.define_builtin('llvm.memset.p0i8.i32', address_dimension_32, 'void')
memset.define_builtin('llvm.memset.p0i8.i64', address_dimension_64, 'void')

memcpy.define_builtin('__memcpy_chk', address_dimension_64, 'void')
msp430_dma_word_copy.define_builtin('__dma_word_copy', address_dimension_32, 'void')