"""Line profiling utility for the TianXuan backend.

Usage in any backend module::

    from analysis.profile_util import profile

    @profile
    def my_function(...):
        ...

When ``TIANXUAN_PROFILE=1`` env var is set AND line_profiler is installed,
the decorator activates real line-level profiling and prints stats at exit.
Otherwise the decorator is a no-op.

To auto-profile ALL public functions in a module automatically, add at the
bottom of the file::

    from analysis.profile_util import auto_profile_module
    auto_profile_module(__name__)

Auto-profiling is also gated by ``TIANXUAN_PROFILE=1``.
"""
import os
import sys
import logging
import atexit

logger = logging.getLogger(__name__)

_PROFILER_ACTIVE = os.environ.get('TIANXUAN_PROFILE', '') == '1'

if _PROFILER_ACTIVE:
    try:
        from line_profiler import profile as _line_profile
    except ImportError:
        logger.warning('TIANXUAN_PROFILE=1 but line_profiler not installed')
        _PROFILER_ACTIVE = False


def profile(func):
    """Decorate a function for line-level profiling.

    Active only when ``TIANXUAN_PROFILE=1`` env var is set and
    line_profiler is installed.
    """
    if _PROFILER_ACTIVE:
        return _line_profile(func)
    return func


def auto_profile_module(module_name: str) -> int:
    """Apply @profile to every public (non-underscore) function in *module_name*.

    Must be called at module level, typically at the bottom of the file::

        from analysis.profile_util import auto_profile_module
        auto_profile_module(__name__)

    Active only when ``TIANXUAN_PROFILE=1`` env var is set and
    line_profiler is installed.  Returns number of functions wrapped, or 0
    when inactive.
    """
    if not _PROFILER_ACTIVE:
        return 0

    try:
        from line_profiler import LineProfiler
    except ImportError:
        return 0

    module = sys.modules.get(module_name)
    if module is None:
        logger.warning('auto_profile_module: module %r not found', module_name)
        return 0

    lp = LineProfiler()
    count = 0
    for name in dir(module):
        if name.startswith('_'):
            continue
        obj = getattr(module, name)
        if callable(obj):
            try:
                lp.add_function(obj)
                count += 1
            except (TypeError, ValueError):
                pass  # skip non-Python or special objects

    atexit.register(lambda: lp.print_stats(output_unit=1e-6))
    logger.info('auto_profile: wrapped %d functions in %s', count, module_name)
    return count
