import asyncio
import os
import threading
import time
import re
from datetime import datetime
from typing import List, Dict, Any
import logging

import pandas as pd

from ivoryos.utils import utils, bo_campaign
from ivoryos.utils.db_models import Script, WorkflowRun, WorkflowStep, db, WorkflowPhase
from ivoryos.utils.global_config import GlobalConfig
from ivoryos.utils.decorators import BUILDING_BLOCKS
from ivoryos.utils.nest_script import validate_and_nest_control_flow

global_config = GlobalConfig()
global deck
deck = None
# global deck, registered_workflows
# deck, registered_workflows = None, None
class HumanInterventionRequired(Exception):
    pass

def pause(reason="Human intervention required"):
    handlers = global_config.notification_handlers
    if handlers:
        for handler in handlers:
            try:
                handler(reason)
            except Exception as e:
                print(f"[notify] handler {handler} failed: {e}")
    # raise error to pause workflow in gui
    raise HumanInterventionRequired(reason)

class ScriptRunner:
    def __init__(self, globals_dict=None):
        self.logger = None
        self.socketio = None
        self.retry = False
        if globals_dict is None:
            globals_dict = globals()
        self.globals_dict = globals_dict
        self.execution_queue = []  # List to hold pending tasks
        self.pause_event = threading.Event()  # A threading event to manage pause/resume
        self.pause_event.set()
        self.stop_pending_event = threading.Event()
        self.stop_current_event = threading.Event()
        self.stop_cleanup_event = threading.Event()
        self.is_running = False
        self.lock = global_config.runner_lock
        self.paused = False
        self.current_app = None
        self.last_progress = 0
        self.last_execution_section = None
        self.waiting_for_input = False
        self.input_value = None
        self.current_task = None

    def handle_input_submission(self, value):
        """Resume execution with user input"""
        if self.waiting_for_input:
            self.input_value = value
            self.pause_event.set()
            return True
        return False

    def toggle_pause(self):
        """Toggles between pausing and resuming the script"""
        self.paused = not self.paused
        if self.pause_event.is_set():
            self.logger.info('Pause script')
            self.pause_event.clear()  # Pause the script
            return "Paused"
        else:
            self.pause_event.set()  # Resume the script
            self.logger.info('Resume script')
            return "Resumed"

    def pause_status(self):
        """Toggles between pausing and resuming the script"""
        return self.paused

    def reset_stop_event(self):
        """Resets the stop event"""
        self.stop_pending_event.clear()
        self.stop_current_event.clear()
        self.stop_cleanup_event.clear()
        self.pause_event.set()

    def get_queue_status(self):
        """Returns the current queue status"""
        queue_status = []
        for i, task in enumerate(self.execution_queue):
            # Basic info
            info = {
                "id": i,
                "name": task.get("run_name", "untitled"),
                "status": "pending",
                "args": f"{task.get('repeat_count', 1)} iteration(s)" if not task.get('config') else f"Config: {len(task.get('config'))} entries"
            }
            
            # Add detailed info for the modal
            details = {}
            if task.get("repeat_count"):
                details["Mode"] = "Repeated Execution"
                details["Repeat Count"] = task.get("repeat_count")
            
            if task.get("config"):
                details["Mode"] = "Configuration List"
                config_data = task.get("config")
                details["Config Entries"] = len(config_data)
                # Show a preview of the first 5 entries to avoid sending too much data
                details["Config Preview (First 5)"] = config_data[:5]
                # If it's a list of dicts, maybe show the first one or just say it's a list
                # For CSV loaded configs, it is a list of dicts
            
            if task.get("batch_size"):
                details["Batch Size"] = task.get("batch_size")
                
            if task.get("history"):
                 details["History File"] = task.get("history")

            # Optimization details
            if task.get("optimizer_cls"):
                details["Mode"] = "Bayesian Optimization"
                details["Optimizer"] = task.get("optimizer_cls").__name__
                
            if task.get("objectives"):
                details["Objectives"] = task.get("objectives")
                
            if task.get("parameters"):
                details["Parameters"] = task.get("parameters")
                
            if task.get("constraints"):
                details["Constraints"] = task.get("constraints")
                
            if task.get("additional_params"):
                details["Additional Params"] = task.get("additional_params")

            info["details"] = details
            
            # Use display_name if available for the main "name" field
            if task.get("display_name"):
                 info["name"] = task.get("display_name")
                 details["Run Name"] = task.get("run_name") # Keep internal name in details
                 
            queue_status.append(info)

        return queue_status

    def get_task_details(self, task_index):
        """Returns the full details for a specific task"""
        try:
            task_index = int(task_index)
            if 0 <= task_index < len(self.execution_queue):
                task = self.execution_queue[task_index]
                details = {}
                if task.get("repeat_count"):
                    details["Mode"] = "Repeated Execution"
                    details["Repeat Count"] = task.get("repeat_count")
                
                if task.get("config"):
                    details["Mode"] = "Configuration List"
                    details["Config Entries"] = len(task.get("config"))
                    # Return full config here
                    details["Full Config"] = task.get("config")
                
                if task.get("batch_size"):
                    details["Batch Size"] = task.get("batch_size")
                    
                if task.get("history"):
                     details["History File"] = task.get("history")

                # Optimization details
                if task.get("optimizer_cls"):
                    details["Mode"] = "Bayesian Optimization"
                    details["Optimizer"] = task.get("optimizer_cls").__name__
                    
                if task.get("objectives"):
                    details["Objectives"] = task.get("objectives")
                    
                if task.get("parameters"):
                    details["Parameters"] = task.get("parameters")
                    
                if task.get("constraints"):
                    details["Constraints"] = task.get("constraints")
                    
                if task.get("additional_params"):
                    details["Additional Params"] = task.get("additional_params")
                    
                return details
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error getting task details: {e}")
        return None

    def remove_task(self, task_index):
        """Removes a task from the queue"""
        try:
            task_index = int(task_index)
            if 0 <= task_index < len(self.execution_queue):
                task = self.execution_queue.pop(task_index)
                if self.logger:
                    self.logger.info(f"Removed task from queue: {task.get('run_name')}")
                return True
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error removing task: {e}")
        return False

    def reorder_tasks(self, task_index, direction):
        """Reorder tasks in the queue"""
        try:
            task_index = int(task_index)
            if direction == "up" and task_index > 0:
                self.execution_queue[task_index], self.execution_queue[task_index - 1] = \
                    self.execution_queue[task_index - 1], self.execution_queue[task_index]
                return True
            elif direction == "down" and task_index < len(self.execution_queue) - 1:
                self.execution_queue[task_index], self.execution_queue[task_index + 1] = \
                    self.execution_queue[task_index + 1], self.execution_queue[task_index]
                return True
        except Exception as e:
            if self.logger:
                self.logger.error(f"Error reordering tasks: {e}")
        return False


    def abort_pending(self):
        """Abort the pending iteration after the current is finished"""
        self.stop_pending_event.set()
        self.logger.info("Abort pending tasks")

    def abort_cleanup(self):
        """Abort the pending iteration after the current is finished"""
        self.stop_cleanup_event.set()
        self.logger.info("Abort cleanup")

    def stop_execution(self):
        """Force stop everything, including ongoing tasks."""
        self.logger.info("Stop execution")
        self.stop_current_event.set()
        self.abort_pending()
        self.abort_cleanup()
        if not self.pause_event.is_set():
            self.pause_event.set()
        if self.lock.locked():
            self.lock.release()


    def run_script(self, script, repeat_count=1, run_name=None, logger=None, socketio=None, config=None,
                   output_path="", compiled=False, current_app=None, history=None, optimizer=None, batch_mode=None,
                   batch_size=1, objectives=None, parameters=None, constraints=None, steps=None, optimizer_cls=None,
                   additional_params=None, on_start=None, display_name=None):


        self.socketio = socketio
        self.logger = logger
        global deck
        if deck is None:
            deck = global_config.deck

        # print("history", history)
        if self.current_app is None:
            self.current_app = current_app
        # time.sleep(1)  # Optional: may help ensure deck readiness

        
        task = {
            "script": script,
            "repeat_count": repeat_count,
            "run_name": run_name,
            "config": config,
            "output_path": output_path,
            "current_app": current_app,
            "compiled": compiled,
            "history": history,
            "optimizer": optimizer,
            "batch_mode": batch_mode,
            "batch_size": batch_size,
            "objectives": objectives,
            "parameters": parameters,
            "constraints": constraints,
            "steps": steps,
            "optimizer_cls": optimizer_cls,
            "additional_params": additional_params,
            "on_start": on_start,
            "display_name": display_name
        }
        
        self.execution_queue.append(task)
        if self.logger:
            self.logger.info(f"Added task to queue: {run_name}")
            
        return self._process_queue()
        
    def _process_queue(self):
        """Process the next task in the queue if the runner is free"""
        # Try to acquire lock without blocking
        if not self.lock.acquire(blocking=False):
            return "queued"

        if not self.execution_queue:
            self.lock.release()
            return "empty"
            
        # Get next task
        task = self.execution_queue.pop(0)
        self.current_task = task # Store current task details
        # todo should check if stop event was set and then check with user to make sure
        #  they want to continue with the queue or abort or pause the queue? in case they need to
        #  manually fix something before continuing
        self.reset_stop_event()

        thread = threading.Thread(
            target=self._run_with_stop_check,
            kwargs=task
        )
        thread.start()
        return thread



    async def exec_steps(self, script, section_name, phase_id, kwargs_list=None, batch_size=1):
        """
        Executes a function defined in a string line by line
        :param func_str: The function as a string
        :param kwargs: Arguments to pass to the function
        :return: The final result of the function execution
        """
        _func_str = script.python_script or script.compile()
        _, return_list = script.config_return()

        global deck
        # global deck, registered_workflows
        if deck is None:
            deck = global_config.deck
        # if registered_workflows is None:
        #     registered_workflows = global_config.registered_workflows

        # for i, line in enumerate(step_list):
        #     if line.startswith("registered_workflows"):
        #
        # func_str = script.compile()
        # Parse function body from string
        temp_connections = global_config.defined_variables
        # Prepare execution environment
        exec_globals = {"deck": deck, "time":time, "pause": pause}  # Add required global objects
        # exec_globals = {"deck": deck, "time": time, "registered_workflows":registered_workflows}  # Add required global objects
        exec_globals.update(temp_connections)

        exec_locals = {}  # Local execution scope

        # Define function arguments manually in exec_locals
        # exec_locals.update(kwargs)
        index = 0
        if kwargs_list:
            results = kwargs_list.copy()
        else:
            results = [{} for _ in range(batch_size)]
        nest_script = validate_and_nest_control_flow(script.script_dict.get(section_name, []))

        await self._execute_steps_batched(nest_script, results, phase_id=phase_id, section_name=section_name)

        return results  # Return the 'results' variable

    def _run_with_stop_check(self, script: Script, repeat_count: int, run_name: str, config,
                             output_path, current_app, compiled, history=None, optimizer=None, batch_mode=None,
                             batch_size=None, objectives=None, parameters=None, constraints=None, steps=None,
                             optimizer_cls=None, additional_params=None, on_start=None, display_name=None):
        if current_app:
            ctx = current_app.app_context()
            ctx.push()

        time.sleep(1)
        
        if on_start:
            try:
                on_start()
            except Exception as e:
                if self.logger:
                    self.logger.error(f"Error in on_start callback: {e}")
        elif self.socketio:
             # Fallback if no callback provided? Or just minimal emit?
             self.socketio.emit('start_task', {'run_name': run_name})

        # _func_str = script.compile()
        # step_list_dict: dict = script.convert_to_lines(_func_str)
        self._emit_progress(1)
        filename = f"{run_name}_{datetime.now().strftime('%Y-%m-%d %H-%M')}.csv"
        error_flag = False
        # create a new run entry in the database
        repeat_mode = "batch" if config else "optimizer" if optimizer else "repeat"
        if optimizer_cls is not None:
            # try:
            if self.logger:
                self.logger.info(f"Initializing optimizer {optimizer_cls.__name__}")
            try:
                optimizer = optimizer_cls(experiment_name=run_name, parameter_space=parameters, objective_config=objectives,
                                      parameter_constraints=constraints, additional_params=additional_params,
                                      optimizer_config=steps, datapath=output_path)
                current_app.config["LAST_OPTIMIZER"] = optimizer
            except Exception as e:
                if self.logger:
                    self.logger.error(f"Error during optimizer initialization: {e.__str__()}")
                self._emit_progress(100)
                if self.lock.locked():
                    self.lock.release()
                return None

        with current_app.app_context():
            run = WorkflowRun(name=script.name or "untitled", platform=script.deck or "deck", start_time=datetime.now(),
                              repeat_mode=repeat_mode
                              )
            db.session.add(run)
            db.session.flush()
            run_id = run.id  # Save the ID
            # overwrite filename with the run specific start time
            filename = f"{run_name}_{run.start_time.strftime('%Y-%m-%d %H-%M-%S')}.csv"
            run.data_path = filename
            db.session.commit()

            # setup run-specific logging to a file using run_id
            log_filename = f"{run_name}_{run.start_time.strftime('%Y-%m-%d %H-%M-%S')}.log"
            log_path = os.path.join(current_app.config["LOG_FOLDER"], log_filename)
            run_file_handler = logging.FileHandler(log_path)
            run_file_handler.setFormatter(logging.Formatter('%(asctime)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S'))
            gui_logger = self.logger
            app_loggers = current_app.config["LOGGERS"]
            app_loggers = app_loggers if isinstance(app_loggers, list) else [app_loggers]
            gui_and_app_loggers = [gui_logger, *app_loggers]
            for logger in gui_and_app_loggers:
                if isinstance(logger, str):
                    logger = logging.getLogger(logger)
                try:
                    logger.addHandler(run_file_handler)
                except Exception as e:
                    self.logger.error(f"Failed to setup logger {logger}: {e}")

            try:
            # if True:
                global_config.runner_status = {"id":run_id, "type": "workflow"}
                # Run "prep" section once
                asyncio.run(self._run_actions(script, section_name="prep", run_id=run_id))
                output_list = []
                _, arg_type = script.config("script")
                _, return_list = script.config_return()
                # Run "script" section multiple times
                if repeat_count:
                    asyncio.run(
                        self._run_repeat_section(repeat_count, arg_type, output_list, script,
                                             run_name, return_list, compiled,
                                             history, output_path, run_id, filename, optimizer=optimizer,
                                             batch_mode=batch_mode, batch_size=batch_size, objectives=objectives)
                    )
                elif config:
                    asyncio.run(
                        self._run_config_section(
                            config, arg_type, output_list, script, run_name, run_id, filename, return_list, output_path,
                            compiled=compiled, batch_mode=batch_mode, batch_size=batch_size
                        )
                    )

                # Run "cleanup" section once
                asyncio.run(self._run_actions(script, section_name="cleanup", run_id=run_id))
                # Reset the running flag when done

            except Exception as e:
                if self.logger:
                    self.logger.error(f"Error during script execution: {e.__str__()}")
                error_flag = True
            finally:
                self._emit_progress(100)
                if self.lock.locked():
                    self.lock.release()
                
                self.current_task = None # Clear current task
                
                # Close run-specific log handler
                for logger in gui_and_app_loggers:
                    if isinstance(logger, str):
                        logger = logging.getLogger(logger)
                    logger.removeHandler(run_file_handler)
                run_file_handler.close()

                # Check for next task in queue
                self._process_queue()


        with current_app.app_context():
            run = db.session.get(WorkflowRun, run_id)
            if run is None:
                if self.logger:
                    self.logger.info("Error: Run not found in database.")
            else:
                run.end_time = datetime.now()
                run.run_error = error_flag
                db.session.commit()


    async def _run_actions(self, script, section_name="", run_id=None):

        if self.logger:
            self.logger.info(f'Executing {section_name} steps')

        # V1.4.8 stop cleanup is optional, credit @Veronica
        if self.stop_cleanup_event.is_set():
            if self.logger:
                self.logger.info(f"Stopping execution during {section_name} section.")
            return None

        phase = WorkflowPhase(
            run_id=run_id,
            name=section_name,
            repeat_index=0,
            start_time=datetime.now()
        )
        db.session.add(phase)
        db.session.flush()
        phase_id = phase.id
        db.session.commit()

        step_outputs = await self.exec_steps(script, section_name, phase_id=phase_id)
        # Save phase-level output
        phase = db.session.get(WorkflowPhase, phase_id)
        phase.outputs = utils.sanitize_for_json(step_outputs)
        phase.end_time = datetime.now()
        db.session.commit()
        return step_outputs

    async def _run_config_section(self, config, arg_type, output_list, script, run_name, run_id, filename, return_list, output_path,
                                  compiled=True, batch_mode=False, batch_size=1):
        if not compiled:
            for i in config:
                try:
                    i = utils.convert_config_type(i, arg_type)
                    compiled = True
                except Exception as e:
                    if self.logger:
                        if isinstance(e, SyntaxError):
                            self.logger.error(f"Error in configuration data: {e.args}")
                            self.logger.error(
                                f"{e.msg} at line {e.lineno}, column {e.offset}: {e.text.strip()}"
                            )
                        else:
                            self.logger.error(e)
                    compiled = False
                    break
        if compiled:
            batch_size = int(batch_size)
            nested_list = [config[i:i + batch_size] for i in range(0, len(config), batch_size)]

            for i, kwargs_list in enumerate(nested_list):
                # kwargs = dict(kwargs)
                if self.stop_pending_event.is_set():
                    if self.logger:
                        self.logger.info(f'Stopping execution during {run_name}: {i + 1}/{len(config)}')
                    break
                if self.logger:
                    self.logger.info(f'Executing {i + 1} of {len(nested_list)} with kwargs = {kwargs_list}')
                progress = ((i + 1) * 100 / len(nested_list)) - 0.1
                self._emit_progress(progress)

                phase = WorkflowPhase(
                    run_id=run_id,
                    name="main",
                    repeat_index=i,
                    parameters=utils.sanitize_for_json(kwargs_list),
                    start_time=datetime.now()
                )
                db.session.add(phase)
                db.session.flush()

                phase_id = phase.id
                db.session.commit()

                output = await self.exec_steps(script, "script", phase_id, kwargs_list=kwargs_list, )
                # print(output)
                phase = db.session.get(WorkflowPhase, phase_id)
                if output:
                    # kwargs.update(output)
                    for output_dict in output:
                        output_list.append(output_dict)
                    phase.outputs = utils.sanitize_for_json(output)
                phase.end_time = datetime.now()
                db.session.commit()

                # save results
                if not script.python_script: # and return_list:
                    if i == 0:
                        self._save_results(filename, arg_type, return_list, output_list, output_path)
                    else:
                        self._save_results_last_row(filename, arg_type, return_list, output_list, output_path)

        return output_list

    async def _run_repeat_section(self, repeat_count, arg_types, output_list, script, run_name, return_list, compiled,
                                  history, output_path, run_id, filename,
                                  optimizer=None, batch_mode=None, batch_size=None, objectives=None):

        if optimizer and history:
            file_path = os.path.join(output_path, history)

            previous_runs = pd.read_csv(file_path)

            expected_cols = list(arg_types.keys()) + list(return_list)

            actual_cols = previous_runs.columns.tolist()

            # NOT okay if it misses columns
            if set(expected_cols) - set(actual_cols):
                if self.logger:
                    self.logger.warning(f"Missing columns from history .csv file. Expecting {expected_cols} but got {actual_cols}")
                raise ValueError("Missing columns from history .csv file.")

            # okay if there is extra columns
            if set(actual_cols) - set(expected_cols):
                if self.logger:
                    self.logger.warning(f"Extra columns from history .csv file. Expecting {expected_cols} but got {actual_cols}")

            optimizer.append_existing_data(previous_runs, file_path)

            for row in previous_runs.to_dict(orient='records'):
                output_list.append(row)



        for i_progress in range(int(repeat_count)):
            if self.stop_pending_event.is_set():
                if self.logger:
                    self.logger.info(f'Stopping execution during {run_name}: {i_progress + 1}/{int(repeat_count)}')
                break

            phase = WorkflowPhase(
                run_id=run_id,
                name="main",
                repeat_index=i_progress,
                start_time=datetime.now()
            )
            db.session.add(phase)
            db.session.flush()
            phase_id = phase.id
            db.session.commit()

            if self.logger:
                self.logger.info(f'Executing {run_name} experiment: {i_progress + 1}/{int(repeat_count)}')
            progress = (i_progress + 1) * 100 / int(repeat_count) - 0.1
            self._emit_progress(progress)

            # Optimizer for UI
            if optimizer:
                try:
                    parameters = optimizer.suggest(n=batch_size)

                    if parameters is None or len(parameters) == 0:
                        self.logger.info("No parameters suggested by optimizer.")
                        raise ValueError("No parameters suggested by optimizer.")

                    if self.logger:
                        self.logger.info(f'Parameters: {parameters}')
                    # Re-fetch phase to update
                    phase = db.session.get(WorkflowPhase, phase_id)
                    phase.parameters = utils.sanitize_for_json(parameters)
                    db.session.commit() # Commit parameters early? Or wait? Let's commit to be safe if exec_steps crashes

                    output = await self.exec_steps(script, "script",  phase_id, kwargs_list=parameters)
                    if output:
                        optimizer.observe(output)
                        
                    else:
                        if self.logger:
                            self.logger.info('No output from script')


                except Exception as e:
                    if self.logger:
                        self.logger.info(f'Optimization error: {e}')
                    break
            else:

                output = await self.exec_steps(script, "script", phase_id, batch_size=batch_size)

            phase = db.session.get(WorkflowPhase, phase_id)
            if output:
                # print("output: ", output)
                output_list.extend(output)
                if self.logger:
                    self.logger.info(f'Output value: {output}')
                phase.outputs = utils.sanitize_for_json(output)

            phase.end_time = datetime.now()
            db.session.commit()

            # save results
            if not script.python_script:  # and return_list:
                if i_progress == 0:
                    self._save_results(filename, arg_types, return_list, output_list, output_path)
                else:
                    self._save_results_last_row(filename, arg_types, return_list, output_list, output_path)

            if optimizer and self._check_early_stop(output, objectives):
                if self.logger:
                    self.logger.info('Early stopping')
                break
                

        return output_list

    def _save_results(self, filename, arg_type, return_list, output_list, output_path):
        """Save the results to the filename"""
        file_path = os.path.join(output_path, filename)
        df = pd.DataFrame(output_list)
        # output_columns = list(arg_type.keys()) + list(return_list)
        # df = df.reindex(columns=output_columns)

        print(f'save df {df} to {file_path}')
        df.to_csv(file_path, index=False)

        if self.logger:
            self.logger.info(f'Results saved to {file_path}')

    def _save_results_last_row(self, filename, arg_type, return_list, output_list, output_path):
        """
        Save the last row to the filename. If the file does not exist, create it with header.
        """
        file_path = os.path.join(output_path, filename)
        df = pd.DataFrame([output_list[-1]])

        df.to_csv(file_path,
                  mode="a",
                  header=not os.path.exists(file_path),
                  index=False,
                  )

        if self.logger:
            self.logger.info(f'Append to results saved to {file_path}')

    def _emit_progress(self, progress):
        self.last_progress = progress
        self.socketio.emit('progress', {'progress': progress})

    def safe_sleep(self, duration: float):
        interval = 1  # check every 1 second
        end_time = time.time() + duration
        while time.time() < end_time:
            if self.stop_current_event.is_set():
                return  # Exit early if stop is requested
            time.sleep(min(interval, end_time - time.time()))

    def get_status(self):
        """Returns current status of the script runner."""
        with self.current_app.app_context():
            return {
                "is_running": self.lock.locked(),
                "paused": self.paused,
                "stop_pending": self.stop_pending_event.is_set(),
                "stop_current": self.stop_current_event.is_set(),
            }


    async def _execute_steps_batched(self, steps: List[Dict], contexts: List[Dict[str, Any]], phase_id, section_name, arg_contexts:List[Dict[str, Any]] = None):
        """
        Execute a list of steps for multiple samples, batching where appropriate.
        """
        for step in steps:
            action = step["action"]
            instrument = step["instrument"]
            action_id = step["id"]
            if action == "if":
                await self._execute_if_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                               section_name=section_name)
            elif action == "repeat":
                await self._execute_repeat_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                                   section_name=section_name)
            elif action == "while":
                await self._execute_while_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                                  section_name=section_name)
            elif instrument == "variable":
                await self._execute_variable_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                                     section_name=section_name)
                # print("Variable executed", "current context", contexts)
            elif instrument == "math_variable":
                await self._execute_variable_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                                     section_name=section_name)
            elif instrument == "input":
                await self._execute_variable_batched(step, contexts, phase_id=phase_id, step_index=action_id,
                                                     section_name=section_name)
            elif instrument == "workflows":
                # Recursively logic for nested workflows
                # print(step.get("workflow", []))
                workflow_steps = validate_and_nest_control_flow(step.get("workflow", []))
                if workflow_steps:
                    if self.socketio:
                        self.socketio.emit('log', {'message': f"Entering workflow: {action} with args: {step.get('args', {})}"})
                    
                    # Inject parameters into context
                    
                    # For batched contexts:
                    workflow_contexts = []
                    for context in contexts:
                        substituted_args =  self._substitute_params( step.get("args", {}), context)
                        # args = step.get("args", {})
                        # for key, value in args.items():
                        #      if isinstance(value, str) and value.startswith("#"):
                        #          context[key] = context.get(value[1:])
                        #      else:
                        #          context[key] = value
                        workflow_contexts.append(substituted_args)
                        # print("context", context)
                        # print("substituted_args", substituted_args)
                    if step.get("batch_action", False):
                        await self._execute_steps_batched(workflow_steps, [contexts[0]], arg_contexts=[workflow_contexts[0]], phase_id=phase_id, section_name=f"{section_name}-{action_id-1}")
                        if len(contexts) > 1:
                            # Propagate any new values from first context to others
                            for key, value in contexts[0].items():
                                for context in contexts[1:]:
                                    if key not in context:
                                        context[key] = value

                    else:
                        for context, workflow_context in zip(contexts, workflow_contexts):
                            # sequentially execute workflow steps
                            await self._execute_steps_batched(workflow_steps, [context], arg_contexts=[workflow_context], phase_id=phase_id,
                                                              section_name=f"{section_name}-{action_id - 1}")

            else:
                # Regular action - check if batch
                if step.get("batch_action", False):
                    # Execute once for all samples
                    consolidate_keys = step.get("consolidate_batch_args", [])
                    if consolidate_keys:
                        # Normalize to list if boolean (backward compat)
                        if isinstance(consolidate_keys, bool):
                            consolidate_keys = [k for k, v in step.get("arg_types", {}).items() if v == "list"]
                        
                        # Aggregate arguments
                        aggregated_args = {}
                        
                        # Initialize based on first context substitution to get keys
                        first_args = self._substitute_params(step["args"], contexts[0])
                        
                        # Identify keys to consolidate
                        # Must be in consolidate_keys AND be a list type (user requirement)
                        target_keys = []
                        arg_types = step.get("arg_types", {})
                        for key in consolidate_keys:
                            if arg_types.get(key) == "list" or key in first_args: # weak check if arg_types invalid
                                target_keys.append(key)
                        
                        # Initialize aggregated args
                        for key, val in first_args.items():
                             if key in target_keys:
                                 aggregated_args[key] = []
                             else:
                                 aggregated_args[key] = val # For non-consolidated args, keep first value

                        # Iterate all contexts to collect lists
                        for ctx in contexts:
                            sub_args = self._substitute_params(step["args"], ctx)
                            for key in target_keys:
                                if key in sub_args:
                                    val = sub_args[key]
                                    if isinstance(val, (list, tuple)):
                                        aggregated_args[key].extend(val)
                                    else:
                                        aggregated_args[key].append(val)
                        
                        await self._execute_action_once(step, contexts[0], arg_contexts=arg_contexts, phase_id=phase_id, step_index=action_id,
                                                            section_name=section_name, override_args=aggregated_args)

                    else:
                        await self._execute_action_once(step, contexts[0], arg_contexts=arg_contexts, phase_id=phase_id, step_index=action_id,
                                                            section_name=section_name)

                else:
                    # Execute for each sample
                    if arg_contexts:
                        for context, arg_context in zip(contexts, arg_contexts):
                            await self._execute_action(step, context, arg_contexts=arg_context, phase_id=phase_id, step_index=action_id,
                                                       section_name=section_name)
                    else:
                        for context in contexts:
                            await self._execute_action(step, context, phase_id=phase_id, step_index=action_id,
                                                       section_name=section_name)
                            self.pause_event.wait()



    async def _execute_if_batched(self, step: Dict, contexts: List[Dict[str, Any]], phase_id, step_index, section_name):
        """Execute if/else block for multiple samples."""
        # Evaluate condition for each sample
        for context in contexts:
            condition = self._evaluate_condition(step["args"]["statement"], context)
            if self.logger:
                self.logger.info(f"Evaluating if {step['args']['statement']}: {condition}")
            if condition:
                await self._execute_steps_batched(step["if_block"], [context], phase_id=phase_id, section_name=section_name)
            else:
                await self._execute_steps_batched(step["else_block"], [context], phase_id=phase_id, section_name=section_name)


    async def _execute_repeat_batched(self, step: Dict, contexts: List[Dict[str, Any]], phase_id, step_index, section_name):
        """Execute repeat block for multiple samples."""
        for context in contexts:
            times = step["args"].get("statement", 1)

            if isinstance(times, str) and times.startswith("#"):
                times = context.get(times[1:])
            # print("repeat times", times, type(times))
            for i in range(times):
                # Add repeat index to all contexts
                # for context in contexts:
                #     context["repeat_index"] = i

                await self._execute_steps_batched(step["repeat_block"], [context], phase_id=phase_id, section_name=section_name)


    async def _execute_while_batched(self, step: Dict, contexts: List[Dict[str, Any]], phase_id, step_index, section_name):
        """Execute while block for multiple samples."""
        max_iterations = step["args"].get("max_iterations", 1000)
        active_contexts = contexts.copy()
        iteration = 0

        while active_contexts and self.stop_current_event.is_set() is False:
            # Filter contexts that still meet the condition
            still_active = []

            for context in active_contexts:
                condition = self._evaluate_condition(step["args"]["statement"], context)
                if self.logger:
                    self.logger.info(f"Evaluating while {step['args']['statement']}: {condition}")
                if condition:
                    context["while_index"] = iteration
                    still_active.append(context)

            if not still_active:
                break

            # Execute for contexts that are still active
            await self._execute_steps_batched(step["while_block"], still_active, phase_id=phase_id, section_name=section_name)
            active_contexts = still_active
            iteration += 1

        # if iteration >= max_iterations:
        #     raise RuntimeError(f"While loop exceeded max iterations ({max_iterations})")

    async def _execute_action(self, step: Dict, context: Dict[str, Any], arg_contexts: Dict[str, Any]=None, phase_id=1, step_index=1, section_name=None, override_args=None):
        """Execute a single action with parameter substitution."""
        # Substitute parameters in args
        result = None
        if self.stop_current_event.is_set():
            return context
        
        if override_args is not None:
            substituted_args = override_args
        elif arg_contexts:
            substituted_args = self._substitute_params(step["args"], arg_contexts)
        else:
            substituted_args = self._substitute_params(step["args"], context)

        # Get the component and method
        instrument = step.get("instrument", "")
        action = step["action"]
        if instrument and "." in instrument:
            instrument_type, instrument = instrument.split(".")
        else:
            instrument_type = ""
        # Execute the action
        while True:
            step_db = WorkflowStep(
                phase_id=phase_id,
                step_index=step_index,
                method_name=action,
                start_time=datetime.now(),
            )
            db.session.add(step_db)
            db.session.flush()
            step_id = step_db.id # Save ID
            db.session.commit() # Commit early to release lock
            
            try:
                # ensure stop pending event is not set before executing action
                if self.stop_pending_event.is_set():
                    # if self.logger:
                        # self.logger.info(f'Stopping execution before executing action {instrument}.{action}')
                    return context
                if self.logger:
                    self.logger.info(f"Executing '{instrument}.{action}' with args {substituted_args}")
                
                section_id = f"{section_name}-{step_index-1}"
                self.last_execution_section = section_id
                self.socketio.emit('execution', {'section': section_id})
                if action == "wait":
                    duration = float(substituted_args["statement"])
                    self.safe_sleep(duration)

                elif action == "pause":
                    msg = substituted_args.get("statement", "")
                    pause(msg)

                elif action == "comment":
                    msg = substituted_args.get("statement", "")
                    if self.logger:
                        self.logger.info(f"Comment: {msg}")

                elif instrument_type == "deck" and hasattr(deck, instrument):
                    component = getattr(deck, instrument)
                    if "_(setter)" in action:
                        action = action.replace("_(setter)", "")
                    if hasattr(component, action):
                        attr = getattr(component, action)

                        if callable(attr):
                            # Execute and handle return value
                            if step.get("coroutine", False):
                                result = await attr(**substituted_args)
                            else:
                                result = attr(**substituted_args)
                        else:
                            # Handle property setter/getter
                            if "value" in substituted_args:
                                setattr(component, action, substituted_args["value"])
                                result = substituted_args["value"]
                            else:
                                result = attr
                        # Store return value if specified
                        # return_var = step.get("return", "")
                        # if return_var:
                        #     context[return_var] = result

                elif instrument_type == "blocks" and instrument in BUILDING_BLOCKS.keys():
                    # Inject all block categories
                    method_collection = BUILDING_BLOCKS[instrument]
                    if action in method_collection.keys():
                        method = method_collection[action]["func"]

                        # Execute and handle return value
                        # print(step.get("coroutine", False))
                        if step.get("coroutine", False):
                            result = await method(**substituted_args)
                        else:
                            result = method(**substituted_args)

                        # # Store return value if specified
                        # return_var = step.get("return", "")
                        # if return_var:
                        #     context[return_var] = result
                else:
                    module = global_config.defined_variables.get(instrument, None)
                    if module is None:
                        raise ValueError(f"Unknown instrument '{instrument}'")
                    method = getattr(module, action)
                    if step.get("coroutine", False):
                        result = await method(**substituted_args)
                    else:
                        result = method(**substituted_args)
                        # Store return value if specified
                return_var = step.get("return", "")
                if return_var and result is not None:
                    result = utils.safe_dump(result)  # sanitizes results to save to be json safe
                    context[return_var] = result
                    if arg_contexts:
                        arg_contexts[return_var] = result

            except HumanInterventionRequired as e:
                self.logger.warning(f"Human intervention required: {e}")
                self.socketio.emit('human_intervention', {'message': str(e)})
                # Instead of auto-resume, explicitly stay paused until user action
                # step.run_error = False
                self.toggle_pause()

            except Exception as e:
                self.logger.error(f"Error during script execution: {e}")
                self.socketio.emit('error', {'message': str(e)})
                
                # Update error status in a fresh transaction
                step_db = db.session.get(WorkflowStep, step_id)
                step_db.run_error = True
                db.session.commit()
                
                self.toggle_pause()
            finally:
                step_db = db.session.get(WorkflowStep, step_id)
                step_db.end_time = datetime.now()
                step_db.output = utils.sanitize_for_json(context) # todo if change so output doesnt include all input values as well then need to update ivoryos.routes.data.data.download_workflow_steps_data_csv
                db.session.commit()

                self.pause_event.wait()
            
            if self.retry:
                # only retry if it errored out
                if step_db.run_error:
                    self.retry = False
                    continue
                self.retry = False
            
            break

        return context

    async def _execute_action_once(self, step: Dict, context: Dict[str, Any], arg_contexts, phase_id, step_index, section_name, override_args=None):
        """Execute a batch action once (not per sample)."""
        # print(f"Executing batch action: {step['action']}")
        return await self._execute_action(step, context, arg_contexts=arg_contexts, phase_id=phase_id, step_index=step_index, section_name=section_name, override_args=override_args)

    @staticmethod
    def _substitute_params(args: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Substitute parameter placeholders like #param_1 with actual values."""
        substituted = {}

        def substitute_vars(value: str, context: Dict[str, Any]) -> Any:
            # Replace placeholders of the form `#var` in a string with values from a context.
            def replacer(match):
                var_name = match.group(1)
                if var_name not in context:
                    raise KeyError(f"Missing context value for '{var_name}'")
                return str(context[var_name])

            return re.sub(r"#(\w+)", replacer, value)

        for key, value in args.items():
            if isinstance(value, str) and value.startswith("#"):
                param_name = value[1:]  # Remove '#'
                substituted[key] = context.get(param_name)
            elif isinstance(value, str):
                # for comment need to substitue #args in the arg with the actual values
                substituted[key] = substitute_vars(value, context)
            else:
                substituted[key] = value

        return substituted

    @staticmethod
    def _evaluate_condition(condition_str: str, context: Dict[str, Any]) -> bool:
        """
        Safely evaluate a condition string with context variables.
        """
        # Create evaluation context with all variables
        eval_context = {}

        # Substitute variables in the condition string
        if isinstance(condition_str, str):
            substituted = condition_str
            for key, value in context.items():
                # Replace #variable with actual variable name for eval
                substituted = substituted.replace(f"#{key}", key)
                # Add variable to eval context
                eval_context[key] = value

            try:
                # Safe evaluation with variables in scope
                result = eval(substituted, {"__builtins__": {}}, eval_context)
                return bool(result)
            except Exception as e:
                raise ValueError(f"Error evaluating condition '{condition_str}': {e}")
        elif isinstance(condition_str, bool):
            return condition_str
        else:
            raise condition_str

    def _check_early_stop(self, output, objectives):
        for row in output:
            all_met = True
            for obj in objectives:
                name = obj['name']
                minimize = obj.get('minimize', True)
                threshold = obj.get('early_stop', None)

                if threshold is None:
                    all_met = False
                    break# Skip if no early stop defined

                value = row[name]
                if minimize and value > threshold:
                    all_met = False
                    break
                elif not minimize and value < threshold:
                    all_met = False
                    break

            if all_met:
                return True  # At least one row meets all early stop thresholds

        return False  # No row met all thresholds

    def prompt_user(self, prompt: str, var_type) -> str:
        result = None
        if self.socketio:
            self.waiting_for_input = True
            self.input_value = None
            self.socketio.emit('request_input', {
                'prompt': prompt,
                'type': var_type
            })
            if self.logger:
                self.logger.info(f"Requesting input: {prompt} ({var_type})")

            # Pause and wait for input
            self.pause_event.clear()
            self.pause_event.wait()

            # Process result
            result = self.input_value
            self.waiting_for_input = False

            # Log result
            if self.logger:
                self.logger.info(f"Input received: {result}")
        else:
            if self.logger:
                self.logger.warning("No socketio connection, skipping input")
        return result

    async def _execute_variable_batched(self, step: Dict, contexts: List[Dict[str, Any]], phase_id, step_index,
                                        section_name):
        """Execute variable assignment for multiple samples."""
        var_name = step["action"]  # "vial" in your example
        var_value = step["args"]["statement"]
        arg_type = step["arg_types"]["statement"]

        for context in contexts:
            if step["instrument"] == "input":
                value = self.prompt_user(var_value, arg_type)
                context[var_name] = value
                continue

            # Substitute any variable references in the value
            if isinstance(var_value, str):
                substituted_value = var_value

                # Replace all variable references (with or without #) with their values
                for key, val in context.items():
                    # Handle both #variable and variable (without #)
                    substituted_value = substituted_value.replace(f"#{key}", str(val))
                    # For expressions like "vial+10", replace variable name directly
                    # Use word boundaries to avoid partial matches
                    import re
                    substituted_value = re.sub(r'\b' + re.escape(key) + r'\b', str(val), substituted_value)

                # Handle based on type
                if arg_type == "float":
                    try:
                        # Evaluate as expression (e.g., "10.0+10" becomes 20.0)
                        result = eval(substituted_value, {"__builtins__": {}}, {})
                        context[var_name] = float(result)
                    except:
                        # If eval fails, try direct conversion
                        context[var_name] = float(substituted_value)

                elif arg_type == "int":
                    try:
                        result = eval(substituted_value, {"__builtins__": {}}, {})
                        context[var_name] = int(result)
                    except:
                        context[var_name] = int(substituted_value)

                elif arg_type == "bool":
                    try:
                        # Evaluate boolean expressions
                        result = eval(substituted_value, {"__builtins__": {}}, {})
                        context[var_name] = bool(result)
                    except:
                        context[var_name] = substituted_value.lower() in ['true', '1', 'yes']

                else:  # "str"
                    # For strings, check if it looks like an expression
                    if any(char in substituted_value for char in ['+', '-', '*', '/', '>', '<', '=', '(', ')']):
                        try:
                            # Try to evaluate as expression
                            result = eval(substituted_value, {"__builtins__": {}}, context)
                            context[var_name] = result
                        except:
                            # If eval fails, store as string
                            context[var_name] = substituted_value
                    else:
                        context[var_name] = substituted_value
            else:
                # Direct numeric or boolean value
                context[var_name] = var_value

