import json
import os
import traceback
from datetime import datetime

from ScEpTIC import llvmir_parser, tools
from ScEpTIC.analysis import get_analysis
from ScEpTIC.analysis.options import AnalysisResultFormat
from ScEpTIC.emulator.intermittent_executor import interruption_managers
from ScEpTIC.emulator.state_retention_manager import StateRetentionManager
from ScEpTIC.emulator.io.input import InputManager
from ScEpTIC.emulator.io.output import OutputManager, OutputSkeleton
from ScEpTIC.emulator.vmstate import VMState
from ScEpTIC.exceptions import ConfigurationException, StopAnomalyFoundException, StopException
from ScEpTIC.AST.transformations.base.instructions.custom_log import CustomLog

from ScEpTIC.emulator.memory.physical.physical_memory_management_unit import PhysicalMemoryManagementUnit


class VM:
    """
    ScEpTIC VM
    """

    def _deprecated_config(self, config):
        self.run_continuous             = config.deprecated.run_continuous
        self.run_locate_memory_test     = config.deprecated.run_locate_memory_test
        self.run_evaluate_memory_test   = config.deprecated.run_evaluate_memory_test
        self.run_input_consistency_test = config.deprecated.run_input_consistency_test
        self.run_output_profiling       = config.deprecated.run_output_profiling_test
        self.run_profiling              = config.deprecated.run_profiling
        self.run_memory_size_measure    = config.deprecated.run_memory_size_measure


    def __init__(self, config):
        self.config = config
        self.parsed = llvmir_parser.parse_file(self.config.program.file, self.config.logging.parser_log_level, self.config.logging.log_section_content)
        self.state = VMState(self, self.config)

        self.state_retention = StateRetentionManager(self.state, self.config.state_retention)
        self.state_retention.process_state_save_routines(self.parsed['function_definitions'], self.parsed['function_declarations'])

        self.state.init_code(self.parsed['function_definitions'], self.parsed['function_declarations'])
        self.state.init_gst(self.parsed['global_vars'])

        if config.result_output.dir_append_datetime:
            current_datetime = datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
            result_dir = f"{config.result_output.test_name}_{current_datetime}"
        else:
            result_dir = f"{config.result_output.test_name}"

        self.result_dir = os.path.join(os.getcwd(), config.result_output.save_directory, result_dir)
        self.save_analysis_output = config.result_output.save_analysis_output

        self.state.stats.test_name = config.result_output.test_name

        self.state.log_dir = os.path.join(self.result_dir, 'logs')

        self._deprecated_config(config)

        self.state.initialized = True

        self.termination_reason = {'event': None, 'event_info': None}

        for mmu in PhysicalMemoryManagementUnit._mmus:
            mmu.init_log_handlers()


    def _save_program_code(self):
        """
        If save_llvmir_code is set, saves the program code into "code" folder inside the save_dir.
        """

        if not self.state.config.result_output.save_code:
            return

        save_dir = os.path.join(self.result_dir, 'code')

        if os.path.exists(save_dir):
            return

        os.makedirs(save_dir, exist_ok=True)

        for function_name in self.state.functions:
            name = function_name.replace(self.state.ir_function_prefix, '')
            file_name = os.path.join(save_dir, f'{name}.txt')

            with open(file_name, 'w') as fp:
                fp.write('{}'.format(self.state.functions[function_name]))


    def _save_vm_state(self, name):
        """
        If save_vm_state is set, saves the simulator state after each test into a states folder.
        """

        if not self.state.config.result_output.save_state:
            return

        save_dir = os.path.join(self.result_dir, 'states')
        os.makedirs(save_dir, exist_ok=True)

        save_file = os.path.join(save_dir, '{}.txt'.format(name))

        with open(save_file, 'w') as fp:
            fp.write(self.get_visual_dump())

        if len(PhysicalMemoryManagementUnit._mmus) == 0:
            return

        mmu_save_dir = os.path.join(save_dir, 'physical_memory')
        os.makedirs(mmu_save_dir, exist_ok=True)

        for memory in PhysicalMemoryManagementUnit._mmus:
            if AnalysisResultFormat.TEXT in self.state.config.analysis.results_formats:
                txt_file = os.path.join(mmu_save_dir, '{}.txt'.format(memory.get_label()))
                with open(txt_file, 'w') as fp:
                    fp.write(memory.get_visual_dump())

            if AnalysisResultFormat.JSON in self.state.config.analysis.results_formats:
                json_file = os.path.join(mmu_save_dir, '{}.json'.format(memory.get_label()))
                with open(json_file, 'w') as fp:
                    json.dump(memory.get_dump(), fp)

    def set_termination_reason(self, event, event_info=None):
        """
        Sets the analysis termination reason
        """
        self.termination_reason = {'event': event, 'event_info': event_info}

    def reset(self):
        """
        Resets the whole status of the simulator.
        """

        self.state_retention.reset()
        self.state.reset()
        self.state.memory.force_nvm_reset()

        self.state.stats.reset()
        self.state.custom_metrics.reset()


    def execute_analysis(self):
        """
        Run all the analysis
        """
        for analysis_name in self.state.config.analysis.enabled_analysis:
            self.run_single_analysis(analysis_name)

        for mmu in PhysicalMemoryManagementUnit._mmus:
            mmu.close_log_handlers()

        if CustomLog.handler is not None:
            CustomLog.handler.close()


    def run_single_analysis(self, analysis_name):
        """
        Executes an analysis and saves its results
        :param analysis_name: analysis name
        """
        # get analysis
        analysis = get_analysis(analysis_name, self)

        self._save_program_code()

        run_exception = None

        print(f"Running {analysis_name} analysis")
        try:
            analysis.run()
            print("Done!")

        except KeyboardInterrupt:
            self.set_termination_reason("KEYBOARD_INTERRUPT")
            print("Analysis stopped. Saving output data...")

        except BaseException as e:
            if not self.state.config.result_output.save_results_after_exception:
                raise e

            self.set_termination_reason("EXCEPTION", {'exception': f"{e}", 'stack_trace': f"{traceback.format_exc()}", 'pc': f"{self.state.register_file.pc}"})
            print("An exception stopped the analysis. Saving output data...")
            print(f"PC: {self.state.register_file.pc}")

            try:
                print(f"Current instruction: {self.state.current_instruction}")

                if self.state.current_instruction.metadata is not None:
                    print(f"  Metadata: {self.state.current_instruction.metadata.retrieve()}")

            except BaseException:
                pass

            run_exception = e

        try:
            self._save_vm_state(analysis_name)
        except BaseException as e:
            print(e)
            print(traceback.format_exc())

        if self.state.config.analysis.save_results:
            print(f"Saving {analysis_name} analysis results")

            save_dir = os.path.join(self.result_dir, 'analysis', analysis_name)

            if not os.path.exists(save_dir):
                os.makedirs(save_dir, exist_ok=True)

            file_base_name = os.path.join(save_dir, analysis_name)

            for result_format in self.state.config.analysis.results_formats:
                analysis.save_result(file_base_name, result_format)
                self.state.custom_metrics.save_metrics(save_dir, result_format)

            with open(os.path.join(save_dir, 'termination_reason.json'), 'w') as fp:
                json.dump(self.termination_reason, fp, indent=4, sort_keys=True)

        if run_exception is not None:
            raise run_exception


    def _save_memory_size(self, memory_size_interruption_manager):
        """
        DEPRECATED
        :param memory_size_interruption_manager:
        :return:
        """

        if memory_size_interruption_manager.__class__.__name__ != 'MemorySizeIdentifier':
            return

        save_dir = os.path.join(self.result_dir, 'memory_size')

        if not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)

        save_file = os.path.join(save_dir, f"results.json")

        with open(save_file, "w") as fp:
            json.dump(self.state.memory_size, fp)


    def _load_interruption_manager(self, module_name, class_name):
        """
        DEPRECATED.
        Loads a given interruption manager module.
        It also enables/disables the controls to be done during runtime operations on data, accordingly
        to the loaded module.
            - For data module enables data anomaly check
            - For input module enables input lookup anomaly check
        """

        interruption_manager = interruption_managers.get_interruption_manager(module_name, class_name)(self.state, self.state_retention)

        if interruption_manager.requires_static_checkpoint and not self.state_retention.static_placement:
            raise ConfigurationException('Unable to run "{}, {}" test: static checkpoint placement required.'.format(module_name, class_name))

        self.state.do_data_anomaly_check = interruption_manager.do_data_anomaly_check
        self.state.input_lookup_enabled = interruption_manager.input_lookup_enabled
        self.state.collect_memory_trace = interruption_manager.collect_memory_trace

        return interruption_manager


    def _get_stop_info(self):
        """
        Returns stop information
        """

        pc_stop_at = ''
        space = '  '
        for pc in self.state.register_file.pc._pc_tracking:
            pc_stop_at += '{}{} -> {}\n'.format(space, pc['function_name'], pc['instruction_number'])
            space += ' '*4

        pc_stop_at += '{}{}\n'.format(space, self.state.register_file.pc)

        return pc_stop_at, self.state.global_clock


    def stop_current_test(self):
        """
        Stops the current test
        """

        self.state.force_stop = True

        pc_stop_at, clock = self._get_stop_info()
        self.state.stats.stop_at(pc_stop_at, clock)


    def run_test(self, module_name, class_name):
        """
        DEPRECATED
        Runs completely a test, given the test module (interruption manager)
        """

        print('Running {} test using {}'.format(module_name, class_name))

        interruption_manager = self._load_interruption_manager(module_name, class_name)

        self._save_program_code()
        self.reset_anomalies()
        self.reset_profiling()
        self.reset()

        while not self.state.program_end_reached:
            try:
                if interruption_manager.intermittent_execution_required():
                    interruption_manager.run_with_intermittent_execution()

                else:
                    self.state.run_step()

            except StopAnomalyFoundException:
                pc_stop_at, clock = self._get_stop_info()
                self.state.stats.stop_at(pc_stop_at, clock, True)
                break

            except StopException:
                pc_stop_at, clock = self._get_stop_info()
                self.state.stats.stop_at(pc_stop_at, clock)
                break

        if self.save_analysis_output:

            if not os.path.exists(self.result_dir):
                os.makedirs(self.result_dir, exist_ok=True)

            anomalies_file_name = os.path.join(self.result_dir, '{}_anomalies.txt'.format(module_name))
            profiling_file_name = os.path.join(self.result_dir, '{}_profiling.txt'.format(module_name))

            self._save_memory_size(interruption_manager)

        else:
            anomalies_file_name = None
            profiling_file_name = None

        self.get_found_anomalies(anomalies_file_name)

        if not self.get_profiling_info(profiling_file_name, module_name):
            self.get_observation_info(profiling_file_name)

        self._save_vm_state(module_name)


    def run_tests(self):
        """
        DEPRECATED
        Runs the tests configured in the config file.
        """

        if self.run_continuous:
            self.run_test('base', 'InterruptionManager')

        if self.run_evaluate_memory_test:
            self.run_test('memory_evaluate', 'EvaluateMemoryAnomaliesInterruptionManager')

        if self.run_locate_memory_test:
            self.run_test('memory_locate', 'LocateMemoryAnomaliesInterruptionManager')

        if self.run_input_consistency_test:
            self.run_test('input', 'InputInterruptionManager')

        if self.run_output_profiling:
            self.run_test('output', 'OutputInterruptionManager')

        if self.run_profiling:
            self.run_test('profiling', 'ProfilingInterruptionManager')

        if self.run_memory_size_measure:
            self.run_test('memory_size', 'MemorySizeIdentifier')


    def get_visual_dump(self):
        """
        Returns a visual dump of the current state of the simulator.
        """

        # registers
        retstr = '[REGISTERS]\n'+str(self.state.register_file.get_visual_dump())
        retstr += "\n"

        # stack
        retstr += '[STACK]\n'+str(self.state.memory.stack)
        retstr += self.state.memory.stack.get_visual_dump()
        retstr += "\n\n"

        # heap
        retstr += '[HEAP]\n'+str(self.state.memory.heap)
        retstr += self.state.memory.heap.get_visual_dump()
        retstr += "\n"

        # gst
        retstr += '[GST]\n'+self.state.memory.gst.get_visual_dump()
        retstr += "\n\n"

        # input table
        retstr += '[INPUT TABLE]\n{}\n\n'.format(tools.fancy_dict_to_str(InputManager.input_table))

        # output table
        retstr += '[OUTPUT TABLE]\n{}\n\n'.format(tools.fancy_dict_to_str(OutputManager.output_table))

        # output idempotency
        retstr += '[OUTPUT IDEMPOTENCY]\n{}\n\n'.format(tools.fancy_dict_to_str(OutputManager.get_measured_idempotency()))

        # global clock
        retstr += 'Global Clock: {}\n\n'.format(self.state.global_clock)

        # stats
        retstr += '[Stats]\n {}\n\n'.format(str(self.state.stats).replace('\n', '\n '))

        return retstr

    def get_found_anomalies(self, outfile = None):
        """
        Prints out the found anomalies. If an outfile is specified, such information is written inside it.
        """

        if len(self.state.anomalies) == 0:
            return None

        retval = "Found anomalies: {}\n\n".format(len(self.state.anomalies))

        for anomaly in self.state.anomalies:
            retval += '{}\n'.format(anomaly)

        if outfile is None:
            print(retval)

        else:
            with open(outfile, 'w') as fp:
                fp.write(retval)


    def get_profiling_info(self, outfile = None, module_name = ''):
        """
        Prints out the gethered profiling information. If an outfile is specified, such information is written inside it.
        """

        if len(self.state.profiling) == 0:
            return False

        if module_name == 'output':
            return self.get_output_profiling_info(outfile)

        retval = 'Observed clock cycles: {}\n\n'.format(len(self.state.profiling))
        spaces = ' ' * 4

        for clock_id in sorted(self.state.profiling):
            checkpoint_pc = self.state.state_save_clock_pc_maps[clock_id]
            retval += 'Checkpoint: {}Global clock: {}\n\n'.format(checkpoint_pc.resolve(), clock_id)

            runs = self.state.profiling[clock_id]

            for run_id in sorted(runs):
                if run_id == 0:
                    retval += '{}Normal execution:\n'.format(spaces)
                else:
                    retval += '{}#{} power failure generated:\n'.format(spaces, run_id)

                run = runs[run_id]

                for io in run:
                    io_str = str(io).replace('\n', '\n{}'.format(spaces*3))
                    retval += '{}{}\n'.format(spaces*2, io_str)

                retval += '\n'

            retval += '\n'

        if outfile is None:
            print(retval)

        else:
            with open(outfile, 'w') as fp:
                fp.write(retval)

        return True


    def get_output_profiling_info(self, outfile = None):
        if len(self.state.profiling) == 0:
            return False

        retval = ''

        spaces = ' ' * 4

        for clock_id in sorted(self.state.profiling):
            checkpoint_pc = self.state.state_save_clock_pc_maps[clock_id]
            retval += 'Checkpoint: {}\n'.format(checkpoint_pc.resolve(), clock_id)

            outs = self.state.profiling[clock_id]

            for out in outs:
                tracking = outs[out]

                if len(tracking) == 0:
                    continue

                retval += '{}Output {}\n'.format(spaces, OutputSkeleton.output_names[out.replace(self.state.ir_function_prefix, '')])


                for track in tracking:
                    out_str = str(track).replace('\n', '\n{}'.format(spaces*3))
                    retval += '{}{}\n'.format(spaces*2, out_str)

                retval += '\n'

            retval += '\n'

        if outfile is None:
            print(retval)

        else:
            with open(outfile, 'w') as fp:
                fp.write(retval)

        return True


    def get_observation_info(self, outfile = None):
        """
        Prints out the gethered profiling information. If an outfile is specified, such information is written inside it.
        """

        if len(self.state.observations) == 0:
            return False

        retval = ''

        for first_checkpoint in self.state.observations:
            retval += 'Checkpoint: {}\n'.format(first_checkpoint.resolve())

            for second_checkpoint in self.state.observations[first_checkpoint]:
                try:
                    second_checkpoint_str = second_checkpoint.resolve()
                except IndexError:
                    second_checkpoint_str = 'Program End\n'

                retval += '{}End checkpoint: {}'.format(' '*4, second_checkpoint_str)
                retval += '{}Input access models:\n{}'.format(' '*6, tools.fancy_dict_to_str(self.state.observations[first_checkpoint][second_checkpoint], 7))

                retval += '\n'

            retval += '\n'

        if outfile is None:
            print(retval)

        else:
            with open(outfile, 'w') as fp:
                fp.write(retval)

        return True


    def reset_anomalies(self):
        """
        Resets the found anomalies.
        """
        self.state.anomalies = []

    def reset_profiling(self):
        """
        Resets the gathered profiling information.
        """
        self.state.profiling = {}