from __future__ import annotations
from contextlib import ContextDecorator, contextmanager
from typing import Any, Generator

from .node_profiler import NodeProfiler, get_current_profiler, set_current_profiler


@contextmanager
def profiler_context() -> Generator[NodeProfiler, None, None]:
    """Context manager that instantiates and activates a NodeProfiler for tree execution."""
    old = get_current_profiler()
    profiler = NodeProfiler()
    set_current_profiler(profiler)
    try:
        yield profiler
    finally:
        set_current_profiler(old)


class NodeProfileExecution(ContextDecorator):
    """Context manager and decorator that profiles a node's execution if a profiler is active,
    or acts as a zero-cost pass-through if no profiler is active."""

    def __init__(self, node: Any):
        self.node = node
        self._ctx = None

    def __enter__(self):
        profiler = get_current_profiler()
        if profiler is not None:
            self._ctx = profiler.profile_node(self.node)
            self._ctx.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._ctx is not None:
            return self._ctx.__exit__(exc_type, exc_val, exc_tb)
        return False


def profile_node_execution(node: Any) -> NodeProfileExecution:
    """Profiles the execution of node if a profiler context is active."""
    return NodeProfileExecution(node)