"""Performance profiling utilities."""

import functools
import logging
import time
from typing import Callable, ParamSpec, TypeVar

import psutil

from abacus_common_logic.utils.formatting import format_bytes

P = ParamSpec('P')
R = TypeVar('R')


def profile(
    func: Callable[P, R] | None = None,
    *,
    logger: logging.Logger | None = None,
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
    """Log execution time and memory usage for a function.

    Can be used as a decorator with or without arguments to profile
    function execution time and memory consumption.

    Args:
        func: The function to be profiled (when used without parentheses).
        logger: Optional logger to use. If not provided, uses a logger
                for this module.

    Returns:
        Decorated function or decorator.

    Examples:
        >>> @profile
        ... def my_func():
        ...     pass

        >>> @profile(logger=custom_logger)
        ... def my_func():
        ...     pass
    """

    def decorator(fn: Callable[P, R]) -> Callable[P, R]:
        """Decorate function to add profiling.

        Args:
            fn: The function to be profiled.

        Returns:
            Wrapped function.
        """
        _logger = logger if logger is not None else logging.getLogger(__name__)
        _process = psutil.Process()

        @functools.wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            """Wrap function for profiling.

            Args:
                *args: Variable length argument list.
                **kwargs: Arbitrary keyword arguments.

            Returns:
                The result of the decorated function.
            """
            mem_before = _process.memory_info().rss
            start_time = time.perf_counter()

            result = fn(*args, **kwargs)

            duration = time.perf_counter() - start_time
            mem_after = _process.memory_info().rss

            delta_str = format_bytes(mem_after - mem_before, show_sign=True)
            mem_total_str = format_bytes(mem_after)

            _logger.info(
                f'Stats: Time: {duration:.2f}s | '
                f'Memory: {mem_total_str} (Delta: {delta_str})'
            )
            return result

        return wrapper

    # Called with arguments: @profile(logger=...)
    if func is None:
        return decorator

    # Called without arguments: @profile
    return decorator(func)