from pathlib import Path

import yappi

from load.projects.cloud.loadtesting import events


class EnableProfiling(events.TimeLimitedEvent):
    def __init__(self, pstats_file_path: Path):
        self._pstats_file_path = pstats_file_path

    def is_in_progress(self):
        return yappi.is_running()

    def start(self, *args, **kwargs):
        yappi.start()

    def stop(self):
        func_stats = yappi.get_func_stats()  # Собираем статистику выполненных функций
        # func_stats.print_all()  # Выведем в консоль собранную статистику

        self._pstats_file_path.parent.mkdir(parents=True, exist_ok=True)

        func_stats.save(self._pstats_file_path, 'PSTAT')  # Сохраняем в формате pstat
        # func_stats.save('/tmp/loadtesting/yappi_output.callgrind', 'CALLGRIND')  # Сохраняем в формате callgrind
        yappi.stop()  # Останавливаем наблюдение
        yappi.clear_stats()  # Очищаем статистику
