import logging

from ScEpTIC import tools
from ScEpTIC.AST.elements.instruction import Instruction
from ScEpTIC.AST.misc.virtual_memory_enum import VirtualMemoryEnum


class AllocaOperation(Instruction):
    """
    AST node of the LLVM Memory Instructions group - Alloca Instruction
    https://llvm.org/docs/LangRef.html#memoryops
    """

    def __init__(self, target, element_type, elements_number, align):
        super().__init__()

        self.target = target
        self.type = element_type
        self.elements_number = int(elements_number)
        self.align = int(align)
        self.is_first = False


    def __str__(self):
        retstr = super().__str__()
        retstr += 'alloca {} x {}'.format(self.type, self.elements_number)
        return retstr

    def run(self, update_program_counter=True):
        """
        Executes the operation and the target assignment.
        (Update program counter ignored -> always True)
        """

        target = self.target.value
        dimension = len(self.type) * self.elements_number

        address = self._vmstate.memory.stack.allocate(dimension, True, self.metadata, self.align)

        # use special write_address (datalayout omitted, value can't be stored in a physical register
        # since it is resolved and directly-placed by the compiler's backend)
        self._vmstate.register_file.write_address(target, address)
        
        # call run's callback
        self._vmstate.on_run(self.tick_count)

        logging.info('[{}] Allocating {} bits in stack at address {} ({}).'.format(self.instruction_type, dimension, address, target))
    

    def get_uses(self):
        """
        Returns a list containing the names of the registers used by this instruction.
        (used by register allocation)
        """

        # no register used
        return []


    # TODO: check if this is necessary (before was returning an empty set)
    #def get_defs(self):
    #    """
    #    Returns a list of registers defined by this instruction.
    #    (used by register allocation)
    #    """
    #
    #        # no register defined
    #        return []

    
    def get_ignore(self):
        """
        Returns a list of register names to be ignored by register allocation.
        For alloca operation it is the target register.
        """

        return self.target.get_uses()


    def get_input_lookup(self):
        """
        Returns the input lookup data for the current operation
        """

        return tools.build_input_lookup_data(None, None)


    def resolve_memory_tag(self, elements):
        """
        Resolves and returns the memory tag of the targeted element
        """
        if self.memory_tag is None:

            if self.metadata is None:
                if self.is_first:
                    self.memory_tag = "first_alloca"
                else:
                    raise Exception(f"No metadata available for {self}")

            else:
                metadata = self.metadata.retrieve()
                self.memory_tag = metadata['variable_name']

        return self.memory_tag


    def resolve_memory_tag_dependency(self, elements):
        if self.memory_tag_dependency is None:
            self.memory_tag_dependency = self.resolve_memory_tag(elements)

        return self.memory_tag_dependency


    def resolve_memory_address_chain(self, elements):
        """
        Returns a list of all the instructions required to get the address of the targeted element(s)
        """
        return [self]


class LoadOperation(Instruction):
    """
    AST node of the LLVM Memory Instructions group - Load Instruction
    https://llvm.org/docs/LangRef.html#memoryops
    """

    def __init__(self, target, load_type, element, align, volatile):
        super().__init__()

        self.target = target
        self.type = load_type
        self.element = element
        self.align = int(align)
        self.is_volatile = volatile

        self.virtual_memory_target = VirtualMemoryEnum.VOLATILE
        self.virtual_memory_normalized = False
        self.has_virtual_memory_copy = False
        self.virtual_memory_copy = None
        self.is_virtual_memory_copy = False

        # In llvmir the arguments of a function are not stored in stack, but passed as virtual registers
        # or immediate values to the call().
        # To emulate the storing of the values onto the stack, save those values as address registers.
        # The register allocation step needs only to set self.is_arg_of_function_call to True, without touching
        # the target virtual register.
        # The arguments are loaded then onto the stack when the function needs them (the store is generated in the llvm ir already)
        self.is_arg_of_function_call = False


    def __str__(self):
        retstr = super().__str__()

        s_type = str(self.type)

        if self.element.type is not None:
            s_type = ''

        retstr += 'load {} {}{}'.format(s_type, self.element, ' [arg]' if self.is_arg_of_function_call else '')

        retstr += f" (${self.virtual_memory_target.value})"
        if self.virtual_memory_normalized:
            retstr += " {NORMALIZED}"

        if self.is_virtual_memory_copy:
            retstr += " {COPY}"

        return retstr
    
    def run(self, update_program_counter=True):
        """
        Executes the load operation and the target assignment.
        (Update program counter ignored -> always True)
        """

        value = self.get_val()
        
        # explained in __init__ comment
        if self.is_arg_of_function_call:
            target = self.target.value
            self._vmstate.register_file.write_address(target, value)
            
            if self._vmstate.input_lookup_enabled:
                input_lookup_data = self.get_input_lookup()
                self._vmstate.register_file.set_address_input_lookup(target, input_lookup_data)

        else:
            self.save_in_target_register(value)

        # call run's callback
        self._vmstate.on_run(self.tick_count)

        logging.info('[{}] Loading {} into {}.'.format(self.instruction_type, value, self.target.value))


    def get_val(self):
        """
        Executes the operation and the target assignment.
        """

        address = self.element.get_val()
        dimension = len(self.type)

        # read from memory
        value = self._vmstate.memory.read(address, dimension)

        return value


    def get_input_lookup(self):
        """
        Returns the input lookup data for the current operation
        """

        address = self.element.get_val()
        
        return self._vmstate.memory.get_cell_input_lookup(address)


    def get_load_address(self):
        """
        Returns the address to be loaded.
        """

        return self.element.get_val()
        

    def replace_reg_name(self, old_reg_name, new_reg_name):
        """
        Replaces the name of a register used by the instruction with a new one.
        (used by register allocation)
        """

        self.element.replace_reg_name(old_reg_name, new_reg_name)
        self.target.replace_reg_name(old_reg_name, new_reg_name)
    
    
    def get_uses(self):
        """
        Returns a list containing the names of the registers used by this instruction.
        (used by register allocation)
        """

        return self.element.get_uses()

    def get_ignore(self):
        """
        Returns a list of register names to be ignored by register allocation.
        For load operation it is the target register, if the load operation is marked to be used
        as argument of a function call, since it will be loaded as a stack offset.
        """

        # if loads an argument of a function call, it is an address register (stack location)
        # so it must be ignored on register allocation.
        if self.is_arg_of_function_call:
            return self.target.get_uses()

        return []


    def resolve_memory_tag(self, elements):
        """
        Resolves and returns the memory tag of the targeted element
        """
        if self.memory_tag is None:
           self.memory_tag = self.element.resolve_memory_tag(elements)

        return self.memory_tag

    def resolve_memory_tag_dependency(self, elements):
        if self.memory_tag_dependency is None:
            self.memory_tag_dependency = self.element.resolve_memory_tag_dependency(elements)

        return self.memory_tag_dependency

    def resolve_memory_address_chain(self, elements):
        """
        Returns a list of all the instructions required to get the address of the targeted element(s)
        """
        return [self, self.element.resolve_memory_address_chain(elements)]

    def get_memory_address(self):
        """
        :return: the address of the targeted memory cell
        """
        return self.get_load_address()


class StoreOperation(Instruction):
    """
    AST node of the LLVM Memory Instructions group - Store Instruction
    https://llvm.org/docs/LangRef.html#memoryops
    """

    def __init__(self, target, value, align, volatile):
        super().__init__()

        self.target = target
        self.value = value
        self.align = int(align)
        self.is_volatile = volatile

        self.virtual_memory_target = VirtualMemoryEnum.VOLATILE
        self.virtual_memory_normalized = False
        self.has_virtual_memory_copy = False
        self.virtual_memory_copy = None
        self.is_virtual_memory_copy = False

        self._omit_target = True


    def __str__(self):
        retstr = super().__str__()

        retstr += 'store {} in {}'.format(self.value, self.target)

        retstr += f" (${self.virtual_memory_target.value})"
        if self.virtual_memory_normalized:
            retstr += " {NORMALIZED}"

        if self.is_virtual_memory_copy:
            retstr += " {COPY}"

        return retstr

    def run(self, update_program_counter=True):
        """
        Executes the operation and the target assignment.
        (Update program counter ignored -> always True)
        """

        address = self.target.get_val()
        dimension = len(self.value.type)
        content = self.value.get_val()

        # write into memory if it is not a dummy write
        if not self.is_part_of_dummy_write:
            self._vmstate.memory.write(address, dimension, content)
        
        if self._vmstate.input_lookup_enabled:
            input_lookup_data = self.value.get_input_lookup()
            self._vmstate.memory.set_cell_input_lookup(address, input_lookup_data)
        
        # call run's callback
        self._vmstate.on_run(self.tick_count)

        logging.info('[{}] Saving {} into {}.'.format(self.instruction_type, content, self.target.value))

    
    def get_uses(self):
        """
        Returns a list containing the names of the registers used by this instruction.
        (used by register allocation)
        """

        return self.value.get_uses() + self.target.get_uses()


    def get_defs(self):
        """
        Returns a list of registers defined by this instruction.
        (used by register allocation)
        """

        return []


    def get_store_address(self):
        """
        Returns the address in which the value will be stored.
        """

        return self.target.get_val()


    def replace_reg_name(self, old_reg_name, new_reg_name):
        """
        Replaces the name of a register used by the instruction with a new one.
        (used by register allocation)
        """

        self.value.replace_reg_name(old_reg_name, new_reg_name)
        self.target.replace_reg_name(old_reg_name, new_reg_name)


    def resolve_memory_tag(self, elements):
        """
        Resolves and returns the memory tag of the targeted element
        """
        if self.memory_tag is None:
            self.memory_tag = self.target.resolve_memory_tag(elements)

        return self.memory_tag

    def resolve_memory_tag_dependency(self, elements):
        if self.memory_tag_dependency is None:
            self.memory_tag_dependency = self.target.resolve_memory_tag_dependency(elements)

        return self.memory_tag_dependency

    def resolve_memory_address_chain(self, elements):
        """
        Returns a list of all the instructions required to get the address of the targeted element(s)
        """
        return [self, self.target.resolve_memory_address_chain(elements)]

    def get_memory_address(self):
        """
        :return: the target address of the store
        """
        return self.get_store_address()


class GetElementPointerOperation(Instruction):
    """
    AST node of the LLVM Memory Instructions group - GetElementPointer Instruction
    https://llvm.org/docs/LangRef.html#memoryops
    """

    def __init__(self, target, element, base_type, indexes, inbounds):
        super().__init__()

        self.target = target
        self.element = element
        self.type = base_type
        # indexes is a list of Values with an additional inrage attribute (either True or False)
        self.indexes = indexes
        self.inbounds = inbounds


    def __str__(self):
        retstr = super().__str__()

        s_indexes = ''

        for i in self.indexes:
            s_indexes = '{}[{}]'.format(s_indexes, i)

        retstr += 'getelementpointer {} {}'.format(self.element, s_indexes)
        return retstr


    def get_val(self):
        """
        Returns the represented absolute address.
        """

        # get relative address to perform computation
        address = self.element.get_val()

        if address in self._vmstate.functions:
            raise NotImplementedError("Direct manipulation of function pointers not supported!")

        prefix, base_address = self._vmstate.memory._parse_absolute_address(address)

        # first index is the offset from the base_address (dimension of spacing = overall size of one element of composition self.type)
        # (is like "pointer" spacing)
        offset = len(self.type) * self.indexes[0].get_val()

        # C-like indexes starts from the second element of self.indexes
        indexes = self.indexes[1:]
        composition = self.type.get_memory_composition()
        #print(f"Index: {indexes}; Composition: {composition}; Initial offset = {offset}; First index: {self.indexes[0]}")

        # each index is an instance of Value
        for index in indexes:
            index = index.get_val()

            composition_elements = composition[0]
            flat = []

            # Array index (composition_elements = number of elements in the array)
            # E.g. [3, 32] or [3, [3, 32]]
            if isinstance(composition_elements, int):
                self.type.flat_composition(composition[1], flat)
                offset += index * sum(flat)
                composition = composition[1]

            # Struct selector (composition is the entire struct data)
            # E.g. [[1, 32], [1, 16], [1, 18]]
            elif isinstance(composition_elements, list):
                prev_elements = composition[:index]
                self.type.flat_composition(prev_elements, flat)
                offset += sum(flat)
                # next element -> sub-struct element
                composition = composition[index]

            else:
                raise NotImplementedError("GetElementPointerOperation() error")

        # offset is calculated in bits, address is in bytes
        if offset % 8 != 0:
            raise ValueError('Offset {} is not a multiple of a byte.'.format(offset))

        #print("-> Offset: {}\n".format(offset))
        offset = offset // 8  # int division

        value = '{}{}'.format(prefix, hex(base_address + offset))

        logging.info('[{}] Address resolved to {}.'.format(self.instruction_type, value))

        return value

    
    def get_uses(self):
        """
        Returns a list containing the names of the registers used by this instruction.
        (used by register allocation)
        """

        uses = self.element.get_uses()
        for index in self.indexes:
            uses = uses + index.get_uses()

        return uses


    def replace_reg_name(self, old_reg_name, new_reg_name):
        """
        Replaces the name of a register used by the instruction with a new one.
        (used by register allocation)
        """

        self.element.replace_reg_name(old_reg_name, new_reg_name)

        for index in self.indexes:
            index.replace_reg_name(old_reg_name, new_reg_name)

        if self.target is not None:
            self.target.replace_reg_name(old_reg_name, new_reg_name)


    def get_input_lookup(self):
        """
        Returns the input lookup data for the current operation
        """

        # element is an address (or a reference to it), so does not have any input lookup info.
        # indexes may have input lookup infos
        lookup = tools.build_input_lookup_data(None, None)

        for index in self.indexes:
            lookup = tools.merge_input_lookup_data(lookup, index.get_input_lookup())

        return lookup


    def resolve_memory_tag(self, elements):
        """
        Resolves and returns the memory tag of the targeted element
        """
        if self.memory_tag is None:
            memory_tag = str(self.element.resolve_memory_tag(elements))

            for index in self.indexes:
                memory_tag += f"[{index.resolve_memory_tag(elements)}]"

            self.memory_tag = memory_tag

        return self.memory_tag

    def resolve_memory_tag_dependency(self, elements):
        if self.memory_tag_dependency is None:
            dep = [None]

            for index in self.indexes:
                dep.append(index.resolve_memory_tag(elements))
                dep.append(index.resolve_memory_tag_dependency(elements))

            self.memory_tag_dependency = dep

        return self.memory_tag_dependency

    def resolve_memory_address_chain(self, elements):
        """
        Returns a list of all the instructions required to get the address of the targeted element(s)
        """
        chain = []

        if self.target is not None:
            chain.append(self)

        chain.append(self.element.resolve_memory_address_chain(elements))

        for index in self.indexes:
            chain.append(index.resolve_memory_address_chain(elements))

        return chain