import io
import sys
import traceback
from damaris.experiments import Experiment


class DurationEstimator:
    """
    Consumes an experiment script's generator and estimates total duration
    without writing any files or invoking the backend.

    Sleeps and synchronizations in the script are tracked instead of
    blocking, so their wall-clock time is included in the estimate.
    """

    def __init__(self):
        pass

    def estimate(self, script):
        """
        Run the experiment script in a sandbox and return a duration estimate.

        :param str script: experiment script source code
        :returns: dict with 'experiment_time', 'sleep_time', 'total_time', and 'job_count'
        :rtype: dict
        """
        sleep_time = [0.0]  # list so it's mutable in nested function

        def tracking_sleep(seconds):
            sleep_time[0] += seconds

        def noop_sync(before=0, waitsteps=0.1):
            pass

        def noop_data():
            return {}

        dataspace = {}
        exp_classes = __import__('damaris.experiments', dataspace, dataspace, ['Experiment'])
        for name in dir(exp_classes):
            if name[:2] == "__" and name[-2:] == "__":
                continue
            dataspace[name] = exp_classes.__dict__[name]
        del exp_classes

        # ranges module for lin_range, log_range, etc.
        ranges = __import__('damaris.tools.ranges', dataspace, dataspace, [])
        for name in dir(ranges):
            if name[:2] == "__" and name[-2:] == "__":
                continue
            dataspace[name] = getattr(ranges, name)
        del ranges

        # inject sandbox helpers (non-blocking versions)
        dataspace["data"] = noop_data
        dataspace["synchronize"] = noop_sync
        dataspace["sleep"] = tracking_sleep

        # execute the script
        try:
            exec(script, dataspace)
        except Exception as e:
            traceback.print_exc()
            raise RuntimeError("Script execution failed") from e

        if "experiment" not in dataspace:
            raise ValueError("Script does not define an experiment() function")

        # consume the generator
        try:
            exp_iterator = dataspace["experiment"]()
        except Exception as e:
            traceback.print_exc()
            raise RuntimeError("Calling experiment() failed") from e

        experiment_time = 0.0
        job_count = 0

        while True:
            try:
                job = next(exp_iterator)
            except StopIteration:
                break
            except Exception as e:
                traceback.print_exc()
                raise RuntimeError("Generator failed") from e

            experiment_time += job.get_length()
            job_count += 1

        total_time = experiment_time + sleep_time[0]

        return {
            "experiment_time": experiment_time,
            "sleep_time": sleep_time[0],
            "total_time": total_time,
            "job_count": job_count,
        }


def estimate_duration(script):
    """
    Convenience function to estimate experiment duration.

    :param str script: experiment script source code
    :returns: dict with 'experiment_time', 'sleep_time', 'total_time', and 'job_count'
    :rtype: dict
    """
    return DurationEstimator().estimate(script)