import copy
import gc
import logging

from ScEpTIC.AST.elements.instructions.other_operations import CallOperation
from ScEpTIC.emulator.io.output import OutputManager
from ScEpTIC.exceptions import MemoryException


class StateRetentionManager:
    """
    Manages checkpoint and restore operations
    """

    def __init__(self, vmstate, config):
        self.vmstate = vmstate

        self.routine_names = {
            'save': config.state_save_function_name,
            'restore': config.state_restore_function_name,
        }

        self.physical_mmu = config.physical_state_retention

        self.volatile_stack = vmstate.memory.memory_positions['stack'] == 'volatile'

        # Ignore state save and restore routines
        CallOperation.ignore.append(self.routine_names['save'])
        CallOperation.ignore.append(self.routine_names['restore'])

        self.use_checkpoints = config.type == 'checkpoint'

        self.static_placement = config.state_save_strategy == 'static_placement' and config.type == 'checkpoint'
        self.probe_energy_buffer = config.state_save_behaviour == 'probe' and config.type == 'checkpoint'
        self.probe_with_adc = config.probe_with_adc and config.type == 'checkpoint'
        self.hibernate_after_save = config.hibernate_after_save and config.type == 'checkpoint'

        self.restore_register_file = config.restore_register_file
        self.restore_stack = config.restore_stack
        self.restore_heap = config.restore_heap
        self.restore_volatile_gst = config.restore_volatile_gst
        self.restore_non_volatile_gst = config.restore_non_volatile_gst
        self.restore_environment = config.restore_environment

        self.calculate_v_save_state = config.calculate_v_save_state
        self.calculate_v_save_state_additional_cc = config.calculate_v_save_state_additional_cc
        self.calculate_v_resume_additional_cc = config.calculate_v_resume_additional_cc
        self.calculate_v_restore_min_cc = config.calculate_v_restore_min_cc
        self.v_save_state = config.v_save_state
        self.v_resume = config.v_resume

        self.callbacks = {
            'before_save': config.before_save_function_name,
            'after_save': config.after_save_function_name,
            'before_restore': config.before_restore_function_name,
            'after_restore': config.after_restore_function_name,
        }

        self.reset()


    def reset(self):
        """
        Resets the saved dump and state
        """
        self.state_dump  = {'register_file': None, 'reg_saving_pool': None, 'stack': None, 'heap': None, 'volatile_gst': None, 'non_volatile_gst': None, 'stack_top_address': None, 'global_clock': 0, 'function_call_lookup': None, 'output': None}
        self.state_saved = {'register_file': None, 'reg_saving_pool': None, 'stack': None, 'heap': None, 'volatile_gst': None, 'non_volatile_gst': None, 'stack_top_address': None, 'global_clock': 0, 'function_call_lookup': None, 'output': None}


    def process_state_save_routines(self, functions, declarations):
        """
        If dynamic checkpoint mechanisms: removes all references to the checkpoint and restore routines (if any).
        If static checkpoint mechanisms: removes all references to restore routine and removes target from checkpoint routine calls.

        NB: Checkpoint routine is considered just as a marker to verify when interrupt and test the execution.
        """

        # remove declaration of state-save routines
        if self.routine_names['save'] in declarations:
            del declarations[self.routine_names['save']]

        # remove declaration of state-restore routines
        if self.routine_names['restore'] in declarations:
            del declarations[self.routine_names['restore']]

        # parse each function and remove restore routine (usually is at the start of the main)
        for name in functions:
            function = functions[name]
            i = 0

            while i < len(function.body):
                instruction = function.body[i]
                
                # if is a restore operation, remove it
                if isinstance(instruction, CallOperation) and instruction.name == self.routine_names['restore']:

                    # preserve label of the instruction to be removed
                    if instruction.label is not None:
                        function.body[i+1].label = copy.deepcopy(instruction.label)

                    del function.body[i]

                # if is a checkpoint operation
                elif isinstance(instruction, CallOperation) and instruction.name == self.routine_names['save']:
                    # if state-save routines not statically placed, remove them
                    if not self.static_placement:
                        # preserve label
                        if instruction.label is not None:
                            function.body[i+1].label = copy.deepcopy(instruction.label)
                        del function.body[i]

                    # if statically placed, remove its target and continue
                    else:
                        function.body[i].target = None
                        i += 1

                else:
                    i += 1

            # update function labels
            function.update_labels()

            # remove from internal calls the state save and restore functions
            # so they are ignored during register allocation
            function._remove_from_calls(self.routine_names['save'])
            function._remove_from_calls(self.routine_names['restore'])

    def execute_state_save(self):
        """
        Saves the state.
        """
        logging.info('[CheckpointManager] Saving checkpoint.')

        # clean
        del self.state_saved['register_file']
        del self.state_saved['reg_saving_pool']
        del self.state_saved['stack']
        del self.state_saved['heap']
        del self.state_saved['volatile_gst']
        del self.state_saved['non_volatile_gst']
        del self.state_saved['stack_top_address']
        del self.state_saved['global_clock']
        del self.state_saved['function_call_lookup']
        del self.state_saved['output']

        gc.collect()

        self._inject_function_call('before_save')

        # save register file
        if self.restore_register_file:
            self.state_saved['register_file'] = self.vmstate.register_file.dump()
            self.state_saved['stack_top_address'] = copy.deepcopy(self.vmstate.memory.stack.top_address)
            self.state_saved['reg_saving_pool'] = copy.deepcopy(self.vmstate.reg_saving_pool)
        else:
            self.state_saved['register_file'] = None
            self.state_saved['stack_top_address'] = None
            self.state_saved['reg_saving_pool'] = None

        # save stack
        if self.restore_stack:
            self.state_saved['stack'] = self.vmstate.memory.stack.dump()
        else:
            self.state_saved['stack'] = None

        # save heap
        if self.restore_heap:
            self.state_saved['heap'] = self.vmstate.memory.heap.dump()
        else:
            self.state_saved['heap'] = None

        # save volatile gst
        if self.restore_volatile_gst:
            volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
            if volatile_gst is not None:
                self.state_saved['volatile_gst'] = volatile_gst.dump()
            else:
                self.state_saved['volatile_gst'] = None
        else:
            self.state_saved['volatile_gst'] = None

        # save non-volatile gst
        if self.restore_non_volatile_gst:
            non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
            if non_volatile_gst is not None:
                self.state_saved['non_volatile_gst'] = non_volatile_gst.dump()
            else:
                self.state_saved['non_volatile_gst'] = None
        else:
            self.state_saved['non_volatile_gst'] = None

        self.state_saved['global_clock'] = self.vmstate.global_clock

        # save function call lookup
        self.state_saved['function_call_lookup'] = copy.deepcopy(self.vmstate.function_call_lookup)

        # output
        if self.restore_environment:
            self.state_saved['output'] = OutputManager.dump()
        else:
            self.state_saved['output'] = None

        if self.physical_mmu is not None:
            self.physical_mmu.write(self.state_saved)

        # stats
        #TODO: new stats
        self.vmstate.stats.checkpoint_executed()

        self._inject_function_call('after_save')


    def execute_state_restore(self):
        """
        Restores the previously saved state.
        """
        logging.info('[CheckpointManager] Restoring checkpoint.')

        self._inject_function_call('before_restore')

        if self.physical_mmu is not None:
            self.state_saved = self.physical_mmu.read()
            #self.physical_mmu.read(self.state_saved)

        # restore stack
        if self.state_saved['stack'] is not None:
            self.vmstate.memory.stack.restore(self.state_saved['stack'])

        # restore register saving pool used by register allocation to track registers
        if self.state_saved['reg_saving_pool'] is not None:
            self.vmstate.reg_saving_pool = copy.deepcopy(self.state_saved['reg_saving_pool'])

        # restore register file
        if self.state_saved['register_file'] is not None:
            self.vmstate.register_file.restore(self.state_saved['register_file'])

            # if stack in volatile memory, the saved state contains only the real portion of the stack
            # so "real" memory zeroing is required.
            if self.volatile_stack:
                # Deallocate and remove memory cells that exceed stack_top_address
                try:
                    self.vmstate.memory.stack._deallocate(self.state_saved['stack_top_address'], False, True)
                except MemoryException as e:
                    print(e)
                    pass

            # if stack not restored, restore ESP (is a relative address of the stack)
            if self.state_saved['stack'] is None:
                # Deallocate memory cells that exceed stack_top_address (just pops)
                try:
                    self.vmstate.memory.stack.deallocate(self.state_saved['stack_top_address'], False, True)
                except MemoryException as e:
                    print(e)
                    pass

        # restore heap
        if self.state_saved['heap'] is not None:
            self.vmstate.memory.heap.restore(self.state_saved['heap'])

        # restore voltaile gst
        if self.state_saved['volatile_gst'] is not None:
            volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
            volatile_gst.restore(self.state_saved['volatile_gst'])

        # restore non-volatile gst
        if self.state_saved['non_volatile_gst'] is not None:
            non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
            non_volatile_gst.restore(self.state_saved['non_volatile_gst'])

        # restore global_clock
        if self.state_saved['global_clock'] > 0:
            self.vmstate.global_clock = self.state_saved['global_clock']

        # restore function call's stack lookup
        if self.state_saved['function_call_lookup'] is not None:
            self.vmstate.function_call_lookup = copy.deepcopy(self.state_saved['function_call_lookup'])

        # restore environment state
        if self.restore_environment:
            OutputManager.restore(self.state_saved['output'])

        # stats
        #TODO: new stats
        self.vmstate.stats.checkpoint_restored()
        
        self._inject_function_call('after_restore')


    def save_dump(self):
        """
        Performs a dump of the vm state
        """

        del self.state_dump['register_file']
        del self.state_dump['reg_saving_pool']
        del self.state_dump['stack']
        del self.state_dump['heap']
        del self.state_dump['volatile_gst']
        del self.state_dump['non_volatile_gst']
        del self.state_dump['global_clock']
        del self.state_dump['function_call_lookup']
        del self.state_dump['output']
        gc.collect()

        # save register file
        self.state_dump['register_file'] = self.vmstate.register_file.dump()
        self.state_dump['reg_saving_pool'] = copy.deepcopy(self.vmstate.reg_saving_pool)

        # save stack
        self.state_dump['stack'] = self.vmstate.memory.stack.dump()

        # save heap
        self.state_dump['heap'] = self.vmstate.memory.heap.dump()

        # save volatile gst
        volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
        if volatile_gst is not None:
            self.state_dump['volatile_gst'] = volatile_gst.dump()
        else:
            self.state_dump['volatile_gst'] = None

        # save non-volatile gst
        non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
        if non_volatile_gst is not None:
            self.state_dump['non_volatile_gst'] = non_volatile_gst.dump()
        else:
            self.state_dump['non_volatile_gst'] = None

        self.state_dump['global_clock'] = self.vmstate.global_clock

        # save function call lookup
        self.state_dump['function_call_lookup'] = copy.deepcopy(self.vmstate.function_call_lookup)

        # save environment state
        self.state_dump['output'] = OutputManager.dump()


    def restore_dump(self):
        """
        Restores the previously taken dump, to make memory consistent again.
        """
        logging.info('[CheckpointManager] Restoring memory state to a consistent one.')

        # restore register file
        self.vmstate.register_file.restore(self.state_dump['register_file'])

        # restore register saving pool used by register allocation to track registers
        self.vmstate.reg_saving_pool = copy.deepcopy(self.state_dump['reg_saving_pool'])

        # restore stack
        self.vmstate.memory.stack.restore(self.state_dump['stack'])

        # restore heap
        self.vmstate.memory.heap.restore(self.state_dump['heap'])

        # restore volatile gst
        volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
        if volatile_gst is not None:
            volatile_gst.restore(self.state_dump['volatile_gst'])

        # restore volatile gst
        non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
        if non_volatile_gst is not None:
            non_volatile_gst.restore(self.state_dump['non_volatile_gst'])

        # restore global_clock, if its dump is present
        self.vmstate.global_clock = self.state_dump['global_clock']

        # restore function call's stack lookup
        self.vmstate.function_call_lookup = copy.deepcopy(self.state_dump['function_call_lookup'])

        # restore output table
        OutputManager.restore(self.state_dump['output'])

        # stats
        # TODO: new stats
        self.vmstate.stats.dump_restored()
        

    def get_vm_state_diff(self):
        """
        Compares the vm state with the vm dump taken during the checkpoint.
        It returns a list containing the memory section names that differs.
        """

        comparations = {
            'register_file': self.vmstate.register_file == self.state_dump['register_file'],
            'stack': self.vmstate.memory.stack == self.state_dump['stack'],
            'heap': self.vmstate.memory.heap == self.state_dump['heap'],
            'volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('volatile') == self.state_dump['volatile_gst'],
            'non_volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile') == self.state_dump['non_volatile_gst'],
            'global_clock': self.vmstate.global_clock == self.state_dump['global_clock'],
            'environment': OutputManager.output_table == self.state_dump['output']
        }

        for name in list(comparations.keys()):
            # remove memory sections names which are not changed.
            if comparations[name]:
                del comparations[name]

        return list(comparations.keys())


    def get_vm_state_diff_elements(self):
        """
        Returns a dictionary containing the differences for each memory component between the current state and the saved dump.
        """

        if self.vmstate.global_clock != self.state_dump['global_clock']:
            global_clock_diff = [{'element': 'global_clock', 'dump_value': self.state_dump['global_clock'], 'current_value': self.vmstate.global_clock}]

        else:
            global_clock_diff = []

        diffs = {
            'register_file': self.vmstate.register_file.diff(self.state_dump['register_file']),
            'stack': self.vmstate.memory.stack.diff(self.state_dump['stack']),
            'heap': self.vmstate.memory.heap.diff(self.state_dump['heap']),
            'volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('volatile').diff(self.state_dump['volatile_gst']),
            'non_volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile').diff(self.state_dump['non_volatile_gst']),
            'global_clock': global_clock_diff,
            'environment': OutputManager.diff(self.state_dump['output'])
        }

        return diffs


    def _inject_function_call(self, name):
        """
        Inject a custom function call
        """
        if name not in self.vmstate.custom_callbacks:
            return

        function_name = self.vmstate.custom_callbacks[name]['name']

        # correct return operation(s) of the injected function
        ret_ops = self.vmstate.custom_callbacks[name]['ret_ops']

        pc = self.vmstate.register_file.pc.save()

        # inject PC
        self.vmstate.register_file.pc.update(function_name, 0)

        while self.vmstate.current_instruction not in ret_ops:
            self.vmstate.run_step()
        
        # restore PC
        self.vmstate.register_file.pc.restore(pc)