from ScEpTIC.AST.misc.virtual_memory_enum import VirtualMemoryEnum
from ScEpTIC.emulator.energy.mcu import MCUClockCycleAction


class SavedStateCalculator:
    """
    Collection of methods that help the calculation of state-saving operations cost.
    """

    def __init__(self, vm):
        """
        :param vm: ScEpTIC vm
        """
        self.vm = vm

        virtual_memory_enabled = 'virtual_memory' in self.vm.state.transformations

        # overwrite functions
        if virtual_memory_enabled:
            self.get_memory_access_mcu_action = self.get_virtual_memory_mcu_action
            self.get_call_ret_mcu_action = self.get_virtual_memory_mcu_action
            self.get_saved_registers = self.virtual_memory_get_saved_registers
            self.get_saved_memory_cells = self.virtual_memory_calculate_saved_memory_cells
        else:
            self.get_saved_memory_cells = self.calculate_saved_memory_cells


    def virtual_memory_get_saved_registers(self, state_save_instr):
        """
        :param state_save_instr: the state saving instruction
        :return: if the state-save operation saves only the program counter and the number of registers
        """
        save_only_pc = not state_save_instr.checkpoint_save_esp
        n_registers = len(state_save_instr.checkpoint_save_regs)

        return save_only_pc, n_registers


    def get_virtual_memory_mcu_action(self, instruction):
        """
        :param instruction: an instruction
        :return: the MCUClockCycleAction of the memory access
        """
        if instruction.virtual_memory_target == VirtualMemoryEnum.NON_VOLATILE:
            return MCUClockCycleAction.NON_VOLATILE_MEMORY_ACCESS, None
        else:
            return MCUClockCycleAction.VOLATILE_MEMORY_ACCESS, None


    def get_memory_access_mcu_action(self, instruction):
        """
        :param instruction: a load or store instruction
        :return: the MCUClockCycleAction of the memory access, the physical memory name (None if not a physical memory access)
        """
        address = instruction.get_memory_address()
        prefix = self.vm.state.memory.extract_prefix_from_address(address)

        # Physical memory access
        if prefix in self.vm.state.memory.mmus:
            return MCUClockCycleAction.PHYSICAL_MEMORY_ACCESS, self.vm.state.memory.mmus[prefix].get_name()

        # Non-volatile memory access
        elif prefix in self.vm.state.memory.nvm_prefixes:
            return MCUClockCycleAction.NON_VOLATILE_MEMORY_ACCESS, None

        # Volatile memory access
        else:
            return MCUClockCycleAction.VOLATILE_MEMORY_ACCESS, None


    def get_saved_registers(self, state_save_instr):
        """
        :param state_save_instr: the state saving instruction
        :return: if the state-save operation saves only the program counter and the number of registers
        """
        save_only_pc = False
        n_registers = self.vm.state.config.register_file.n_physical_registers

        return save_only_pc, n_registers


    def get_call_ret_mcu_action(self, instruction):
        """
        :param instruction: a call or ret instruction
        :return: the MCUClockCycleAction of the instruction memory accesses, the physical memory name (None if not a physical memory access)
        """
        # Physical memory access
        if self.vm.state.memory.mmu_targets['stack'] is not None:
            return MCUClockCycleAction.PHYSICAL_MEMORY_ACCESS, self.vm.state.memory.mmu_targets['stack'].get_name()
        # Non-volatile memory access
        if self.vm.state.memory.nvm_content['stack']:
            return MCUClockCycleAction.NON_VOLATILE_MEMORY_ACCESS, None

        # Volatile memory access
        else:
            return MCUClockCycleAction.VOLATILE_MEMORY_ACCESS, None


    def virtual_memory_calculate_saved_memory_cells(self, state_save_instr):
        """
        :param state_save_instr: the state saving instruction
        :return: the number of memory cells that need to be preserved
        """
        if state_save_instr.checkpoint_save_ram:
            return self.calculate_saved_memory_cells(state_save_instr)

        return 0


    def calculate_saved_memory_cells(self, state_save_instr):
        """
        :param state_save_instr: the state saving instruction
        :return: the number of memory cells that need to be preserved
        """
        ignore_str_bss = '@.str'

        n_global_vars_memory_cells = 0
        n_stack_memory_cells = 0
        n_heap_memory_cells = 0

        # Global variables
        if self.vm.state_retention.restore_volatile_gst:
            gst = self.vm.state.memory.gst._get_gst_from_ram_name('volatile')
            # gst may be disabled
            if gst is not None:
                # Real cells -> the ones without @.str
                gst_cells = list(filter(lambda x: x.metadata not in ignore_str_bss, gst._memory.values()))
                n_global_vars_memory_cells = len(gst_cells)

                #for memory_cell in gst._memory.values():
                #    if ignore_str_bss not in memory_cell.metadata:
                #        n_global_vars_memory_cells += 1

        # Stack
        if self.vm.state_retention.restore_stack and self.vm.state.memory.memory_positions['stack'] == 'volatile':
            stack = self.vm.state.memory.stack

            # Active cells -> the ones whose address is not exceding stack.top_address
            active_cells = list(filter(lambda x: x < stack.top_address, stack._memory.keys()))
            n_stack_memory_cells = len(active_cells)

            #for address, memory_cell in sorted(stack._memory.items()):
            #    if address < stack_pointer:
            #        n_stack_memory_cells += 1

        # Heap
        if self.vm.state_retention.restore_heap and self.vm.state.memory.memory_positions['heap'] == 'volatile':
            heap = self.vm.state.memory.heap

            # Active cells -> the ones not marked as garbage
            active_cells = list(filter(lambda x: not x.garbage, heap._memory.values()))
            n_heap_memory_cells = len(active_cells)

            #for memory_cell in heap._memory.values():
            #    if not memory_cell.garbage:
            #        n_heap_memory_cells += 1

        return n_global_vars_memory_cells + n_stack_memory_cells + n_heap_memory_cells