"""
Profile
=================

This module provides utilities for profiling and tracing execution of code.

It includes:

- **profile_test**: A context manager for profiling code execution using the
  built-in `cProfile` module. The profiler can dump the stats in various formats:
  - `.dat`: Raw profile dump.
  - `.txt`: Human-readable text output.
  - `.png`: Graphical representation using `gprof2dot` and Graphviz's `dot` tool.

- **trace_line**: A decorator that launches a background thread to trace and
  print the current line number being executed at regular intervals. This is
  especially useful when debugging long-running processes that might be
  terminated by the operating system without a traceback.

Usage Examples
--------------

Profiling Example:

.. code-block:: python

    with profile_test(enable=True, output_file='profile.txt'):
        # Code to profile
        do_something()

Trace Line Example:

.. code-block:: python

    @trace_line(interval=0.1)
    def long_running_function():
        # Some long-running process
        process_data()

    long_running_function()

"""
import functools
import sys
import threading
import time

from neuralib.util.verbose import fprint

__all__ = ['profile_test', 'trace_line']


class profile_test:
    def __init__(self, enable=False, output_file='profile.png'):
        self._enable = enable
        self._profile = None
        self._output_file = output_file

    def __enter__(self):
        if self._enable:
            import cProfile

            self._profile = cProfile.Profile()
            self._profile.enable()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._enable:
            self._profile.disable()

            if self._output_file.endswith('.dat'):
                self._profile.dump_stats(self._output_file)

            elif self._output_file.endswith('.txt'):
                import pstats

                stat = pstats.Stats(self._profile)

                with open(self._output_file, 'w') as f:
                    stat.stream = f
                    stat.print_stats()

            elif self._output_file.endswith('.png'):
                import subprocess

                f1 = self._output_file.replace('.png', '.dat')
                self._profile.dump_stats(f1)
                cmd_line = f'python -m gprof2dot -f pstats {f1} | dot -T png -o {self._output_file}'

                with subprocess.Popen(['bash', '-c', cmd_line]) as proc:
                    proc.wait()


def trace_line(f=None, *, interval: float = 0.1):
    """
    Decorator for trace line number in real time using threading.

    **Use Case**

    - OS directly kill job without traceback

    :param f: Decorated function
    :param interval: Time interval for trace (in second).
    """

    def _decorator(f):

        @functools.wraps(f)
        def _wrapper(*args, **kwargs):
            done = False
            _file = None
            _lineno = None
            _prev_time = None

            def update():
                nonlocal done
                nonlocal _file
                nonlocal _lineno
                nonlocal _prev_time

                while not done:
                    time.sleep(interval)

                    # noinspection PyUnresolvedReferences,PyProtectedMember
                    info = sys._current_frames()[caller_thread]

                    if info.f_code.co_filename != _file or info.f_lineno != _lineno:

                        current_time = time.time()
                        if _prev_time is not None:
                            elapsed_time = current_time - _prev_time
                        else:
                            elapsed_time = 0  # First iteration

                        _file = info.f_code.co_filename
                        _lineno = info.f_lineno
                        _prev_time = current_time

                        fprint(f'File: {_file} in line {_lineno} - Time elapsed: {elapsed_time:.3f} sec',
                               vtype='debug',
                               flush=True)

            thread = threading.Thread(target=update)
            caller_thread = threading.current_thread().ident
            thread.start()

            try:
                result = f(*args, **kwargs)
            finally:
                done = True  # Signal to stop the memory monitoring thread
                thread.join()  # Ensure the thread finishes

            return result

        return _wrapper

    if f is None:
        return _decorator
    else:
        return _decorator(f)
