#!/usr/bin/env python3
'''
       _/      _/                                                _/
      _/_/    _/    _/_/    _/_/_/  _/_/      _/_/      _/_/_/        _/_/_/
     _/  _/  _/  _/_/_/_/  _/    _/    _/  _/_/_/_/  _/_/      _/  _/_/
    _/    _/_/  _/        _/    _/    _/  _/            _/_/  _/      _/_/
   _/      _/    _/_/_/  _/    _/    _/    _/_/_/  _/_/_/    _/  _/_/_/


Copyright 2025 bdunahu

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Commentary:
Code:
'''

from causal_event_loop import CausalEventLoop
from collections import defaultdict
from html_gen import plot_results
import argparse
import asyncio
import os
import random
import signal
import sys
import threading
import time
import traceback
import types

class Experiment:

    # event loops participating in this this experiment
    _loops = []
    # a key-value pair where keys represent a handle and values
    # represent the period of time waiting
    samples = None

    def __init__(self, loops):
        self._loops = loops
        self.samples = defaultdict(lambda: 0)

    def get_loops(self):
        return [l for l in self._loops if l.is_running()]

    def add_handles(self, handles, loop, time):
        for h in handles:
            h.append(loop._thread_id)
            self.samples[tuple(h)] += time

class Nemesis(object):

    # the (ideal) interval between samples
    signal_interval = 0.0
    # the timestamp which the last sample was taken
    last_sample = None

    # the current experiment being run
    experiment_data = None
    # the coroutine the current experiment is speeding up
    experiment_coro = None
    # the speedup of the current experiment
    experiment_spdp = None

    # results from previous experiments. Keys represent names of coroutines.
    results = defaultdict(lambda: defaultdict(lambda: []))

    # the file to write results to
    filename = None

    # The base duration of each performance experiment
    e_duration = None

    # The number of seconds remaining in this performance experiment.
    r_duration = None

    # A mapping of event loops to the previous running coroutine.
    prev_coro = defaultdict(lambda: None)

    @staticmethod
    def __init__(e_duration, filename, signal_interval=0.01):
        Nemesis.signal_interval = signal_interval
        Nemesis.e_duration = e_duration
        Nemesis.filename = filename
        Nemesis.r_duration = 0

    @staticmethod
    def start():
        Nemesis.last_sample = time.perf_counter()
        signal.signal(signal.SIGALRM,
                      Nemesis._signal_handler)
        signal.setitimer(signal.ITIMER_REAL,
                         Nemesis.signal_interval,
                         Nemesis.signal_interval)

    @staticmethod
    def stop():
        signal.setitimer(signal.ITIMER_REAL, 0)
        plot_results(Nemesis.results, Nemesis.filename)
        print(f"Wrote {Nemesis.filename}")

    @staticmethod
    def _start_experiment(coro, speedup):
        Nemesis.r_duration = Nemesis.e_duration * speedup
        Nemesis.prev_coro = defaultdict(lambda: None)
        Nemesis.experiment_coro = coro
        Nemesis.experiment_spdp = speedup

        loops = Nemesis._get_event_loops()
        for l in loops:
            if not isinstance(l, CausalEventLoop):
                raise RuntimeException("Nemesis requires a custom event loop to insert slowdowns. It does not work on programs which change the event loop policy.")
            l._set_dilation(1.0 / speedup)

        Nemesis.experiment_data = Experiment(loops)

    @staticmethod
    def _stop_experiment():
        if Nemesis.experiment_data is not None:
            print(f'finished running experiment on {Nemesis.experiment_coro}')
            results = Nemesis.experiment_data.samples
            Nemesis.results[Nemesis.experiment_coro][Nemesis.experiment_spdp].append(results)
            del Nemesis.experiment_data

    @staticmethod
    def _signal_handler(sig, frame):
        curr_sample = time.perf_counter()
        passed_time = curr_sample - Nemesis.last_sample
        Nemesis.last_sample = curr_sample


        if getattr(Nemesis, 'experiment_data', None):
            loops = Nemesis.experiment_data.get_loops()
            exp_coro = Nemesis.experiment_coro
            for loop in loops:
                coro = Nemesis._get_current_coro(loop)
                prev_coro = Nemesis.prev_coro[loop]
                if not prev_coro == coro:
                    if prev_coro == exp_coro:
                        loop.ping_exit_coro()
                    elif coro == exp_coro:
                        loop.ping_enter_coro()
                    Nemesis.prev_coro[loop] = coro

                loop._update_ready(True)
                handles = Nemesis._get_waiting_handles(loop)
                Nemesis.experiment_data.add_handles(handles, loop, passed_time)

            Nemesis.r_duration -= passed_time
            if (Nemesis.r_duration <= 0):
                Nemesis._stop_experiment()

        else:
            coros = []
            loops = Nemesis._get_event_loops()
            for loop in loops:
                coro = Nemesis._get_current_coro(loop)
                if coro is not None:
                    coros.append(coro)
                if coros:
                    Nemesis._start_experiment(random.choice(coros),
                                              Nemesis._select_speedup())

    @staticmethod
    def _get_waiting_handles(loop):
        handles = []
        for handle in loop._ready:
            handle_info = Nemesis._parse_handle(handle)
            handles.append(handle_info)
        return handles

    @staticmethod
    def _parse_handle(handle):
        cb = handle._callback
        if isinstance(getattr(cb, '__self__', None), asyncio.tasks.Task):
            task = cb.__self__
            coro = task.get_coro()
            return [task.get_name(), Nemesis._get_coro_name(coro)]
        else:
            return [str(type(handle).__name__), cb.__name__]

    @staticmethod
    def _get_current_coro(loop):
        tid = loop._thread_id
        assert tid, f"{loop} is not running, yet we attempted to sample it!"

        frame = sys._current_frames().get(tid)
        fname = frame.f_code.co_filename
        while not Nemesis._should_trace(fname):
            if frame:
                frame = frame.f_back
            else:
                break
            if frame:
                fname = frame.f_code.co_filename
        if frame and frame.f_generator:
            return frame.f_generator.cr_code.co_name
        return None

    @staticmethod
    def _get_event_loops():
        '''Returns each thread's event loop, if it exists.'''
        loops = []
        for t in threading.enumerate():
            frame = sys._current_frames().get(t.ident)
            if frame:
                loop = Nemesis._walk_back_until_loop(frame)
                if loop and loop not in loops:
                    loops.append(loop)
        return loops

    @staticmethod
    def _walk_back_until_loop(frame):
        '''Walks back the callstack until we are in a method named '_run_once'.
        If this is ever true, we assume we are in an Asyncio event loop method,
        and check to see if the 'self' variable is indeed and instance of
        AbstractEventLoop. Return this variable if true.'''
        while frame:
            if frame.f_code.co_name == '_run_once' and 'self' in frame.f_locals:
                loop = frame.f_locals['self']
                if isinstance(loop, asyncio.AbstractEventLoop):
                    return loop
            else:
                frame = frame.f_back
        return None

    @staticmethod
    def _should_trace(filename):
        '''Returns FALSE if filename is uninteresting to the user.
        Don't depend on this. It kind of sucks.'''
        if not filename:
            return False
        if '/gnu/store' in filename:
            return False
        if '/usr/local/lib/python' in filename:
            return False
        if 'site-packages' in filename:
            return False
        if 'propcache' in filename:
            return False
        if '.pyx' in filename:
            return False
        if filename[0] == '<':
            return False
        if 'nemesis' in filename:
            return False
        return True

    def _select_speedup():
        '''
        Returns a random speedup between 0% to 100%, in multiples of 5%.
        Because a baseline is needed to calculate effect on program
        performance, selects a speedup of 0 with 50% probability.
        '''
        r1 = random.randint(-9, 20)
        return 1 + max(0, r1) / 20

    @staticmethod
    def _get_coro_name(coro):
        '''
        Stolen from _format_coroutine in cpython/Lib/asyncio/coroutines.py
        '''
        # Coroutines compiled with Cython sometimes don't have
        # proper __qualname__ or __name__.  While that is a bug
        # in Cython, asyncio shouldn't crash with an AttributeError
        # in its __repr__ functions.
        if hasattr(coro, '__qualname__') and coro.__qualname__:
            coro_name = coro.__qualname__
        elif hasattr(coro, '__name__') and coro.__name__:
            coro_name = coro.__name__
        else:
            # Stop masking Cython bugs, expose them in a friendly way.
            coro_name = f'<{type(coro).__name__} without __name__>'
        return f'{coro_name}()'

the_globals = {
    '__name__': '__main__',
    '__doc__': None,
    '__package__': None,
    '__loader__': globals()['__loader__'],
    '__spec__': None,
    '__annotations__': {},
    '__builtins__': globals()['__builtins__'],
    '__file__': None,
    '__cached__': None,
}


if __name__ == "__main__":
    # parses CLI arguments and facilitates profiler runtime.
    parser = argparse.ArgumentParser(
        usage='%(prog)s [args] -- prog'
    )

    parser.add_argument('-i', '--interval',
                        help='The minimum amount of time inbetween \
                        samples in seconds.',
                        metavar='',
                        type=float,
                        default=0.01)
    parser.add_argument('-e', '--experiment-duration',
                        help='The performance experiment duration. Defaults to 3 seconds.',
                        metavar='',
                        type=float,
                        default=3)
    parser.add_argument('-f', '--filename',
                        help='The filename to write results to.',
                        metavar='',
                        type=str,
                        default="results.html")
    parser.add_argument('prog',
                        type=str,
                        nargs='*',
                        help='Path to the python script and its arguments.')
    args = parser.parse_args()

    sys.argv = args.prog
    try:
        with open(args.prog[0], 'r', encoding='utf-8') as fp:
            code = compile(fp.read(), args.prog[0], "exec")
            Nemesis(args.experiment_duration,
                    args.filename,
                    args.interval).start()
            exec(code, the_globals)
            Nemesis.stop()
    except Exception:
        traceback.print_exc()