import os
import csv
import threading
import time
from datetime import datetime

from ivoryos.utils import utils
from ivoryos.utils.db_models import Script


class ScriptRunner:
    def __init__(self, globals_dict=None):
        if globals_dict is None:
            globals_dict = {}
        self.globals_dict = globals_dict

        self.stop_event = threading.Event()
        self.is_running = False
        self.lock = threading.Lock()

    def reset_stop_event(self):
        self.stop_event.clear()

    def stop_execution(self):
        self.stop_event.set()
        # print("Stop pending tasks")

    def run_script(self, script, repeat_count=1, run_name=None, logger=None, socketio=None, config=None, bo_args=None,
                   output_path=""):
        time.sleep(1)
        with self.lock:
            if self.is_running:
                logger.info("System is busy. Please wait for it to finish or stop it before starting a new one.")
                return None
            self.is_running = True

        self.reset_stop_event()

        thread = threading.Thread(target=self._run_with_stop_check,
                                  args=(script, repeat_count, run_name, logger, socketio, config, bo_args, output_path))
        thread.start()
        return thread

    def _run_with_stop_check(self, script: Script, repeat_count, run_name, logger, socketio, config, bo_args,
                             output_path):
        time.sleep(1)
        try:
            # Run "prep" section once
            script_dict = script.script_dict
            self._run_actions(script_dict.get("prep", []), section_name="prep", run_name=run_name, logger=logger)
            output_list = []
            _, arg_type = script.config("script")
            _, return_list = script.config_return()

            # Run "script" section multiple times
            if repeat_count:
                self._run_repeat_section(repeat_count, bo_args, output_list, return_list, run_name, logger, socketio)
            elif config:
                self._run_config_section(config, arg_type, output_list, return_list, run_name, logger, socketio)

            # Run "cleanup" section once

            self._run_actions(script_dict.get("cleanup", []), section_name="cleanup", run_name=run_name, logger=logger)

            # Save results if necessary
            if output_list:
                self._save_results(run_name, arg_type, return_list, output_list, logger, output_path)
        except Exception as e:
            logger.error(f"Error during script execution: {e}")
        finally:
            with self.lock:
                self.is_running = False  # Reset the running flag when done

    def _run_actions(self, actions, section_name="", run_name=None, logger=None):
        for action in actions:
            if self.stop_event.is_set():
                logger.info(f"Stopping execution during {section_name} section.")
                break
            logger.info(f'Executing {section_name} steps')
            fname = f"{run_name}_{section_name}"
            function = self.globals_dict[fname]
            function()

    def _run_config_section(self, config, arg_type, output_list, return_list, run_name, logger, socketio):
        compiled = True
        for i in config:
            try:
                i = utils.convert_config_type(i, arg_type)
            except Exception as e:
                logger.info(e)
                compiled = False
                break
        if compiled:
            for i, kwargs in enumerate(config):
                kwargs = dict(kwargs)
                logger.info(f'Executing {i + 1} of {len(config)} with kwargs = {kwargs}')
                progress = (i + 1) * 100 / len(config)
                socketio.emit('progress', {'progress': progress})
                fname = f"{run_name}_script"
                function = self.globals_dict[fname]
                output = function(**kwargs)
                if output:
                    kwargs.update(output)
                    output_list.append(kwargs)

    def _run_repeat_section(self, repeat_count, bo_args, output_list, return_list, run_name, logger, socketio):
        if bo_args:
            logger.info('Initializing optimizer...')
            ax_client = utils.ax_initiation(bo_args)
        for i in range(int(repeat_count)):
            if self.stop_event.is_set():
                logger.info(f'Stopping execution during {run_name}: {i + 1}/{int(repeat_count)}')
                break
            logger.info(f'Executing {run_name}: {i + 1}/{int(repeat_count)}')
            progress = (i + 1) * 100 / int(repeat_count)
            socketio.emit('progress', {'progress': progress})
            if bo_args:
                try:
                    parameters, trial_index = ax_client.get_next_trial()
                    logger.info(f'Output value: {parameters}')
                    fname = f"{run_name}_script"
                    function = self.globals_dict[fname]
                    output = function(**parameters)
                    # output = eval(f"{run_name}_script(**{parameters})")
                    ax_client.complete_trial(trial_index=trial_index, raw_data=output)
                except Exception as e:
                    logger.info(f'Optimization error: {e}')
                    break
            else:
                fname = f"{run_name}_script"
                function = self.globals_dict[fname]
                output = function()

            if output:
                output_list.append(output)
                logger.info(f'Output value: {output}')

    @staticmethod
    def _save_results(run_name, arg_type, return_list, output_list, logger, output_path):
        if arg_type:
            args = list(arg_type.keys())
            args.extend(return_list)
            filename = run_name + "_" + datetime.now().strftime("%Y-%m-%d %H-%M") + ".csv"
            file_path = os.path.join(output_path, filename)
            with open(file_path, "w", newline='') as file:
                writer = csv.DictWriter(file, fieldnames=args)
                writer.writeheader()
                writer.writerows(output_list)
            logger.info(f'Results saved to {file_path}')
