from contextlib import contextmanager
from typing import Generic, TypeVar
import threading

from soniclue_basekit.exec_time_profiler import ExecTimeProfiler, nop_profiler
from soniclue_basekit.pytest_utils import is_pytest_env


_T = TypeVar("_T")


class LockProtector(Generic[_T]):
    """Wraps an object with lock protection.

    Notes:
        1. This class also supports lock acquisition order by adding references to "inner" protectors.
        2. We will check that these inner ones should not be locked by the same thread prior to
           this protector.
        3. No circular inner relationship allowed and the check is recursive.
    """

    def __init__(
        self,
        obj: _T,
        profiler: ExecTimeProfiler = nop_profiler,
        check_inner_only_in_pytest: bool = False,
    ):
        """Constructor.

        Args:
            obj: Object to be lock protected.
            profiler: Execution time profiler; default to dummy profiler.
            check_inner_only_in_pytest: Performs inner lock check only in pytest.
        """
        assert obj is not None
        self._object = obj
        self._lock = threading.Lock()
        self._owner_thread: threading.Thread | None = None
        self._inner_protectors: set[LockProtector] = set()
        self._profiler = profiler
        self._check_inner_only_in_pytest = check_inner_only_in_pytest

    def add_inner_protector(self, inner: "LockProtector") -> None:
        """Adds reference to an "inner" lock."""
        assert isinstance(inner, LockProtector)
        self._inner_protectors.add(inner)

    def _assert_acquisition_order(self, visited: set | None = None) -> None:
        """Raises if any inner protector is already held by the current thread.

        Args:
            visited: Tracks already-visited nodes to detect circular relationships.
        """
        if visited is None:
            visited = set()
        if self in visited:
            raise RuntimeError("Circular inner-protector relationship detected")
        visited.add(self)
        current = threading.current_thread()
        for inner in self._inner_protectors:
            inner._assert_acquisition_order(visited)
            if inner._lock.locked() and inner._owner_thread is current:
                raise RuntimeError(f"Inner protector already locked by current thread: {inner}")
        visited.discard(self)

    @contextmanager
    def lock(self, caller: str = "caller"):
        """Yields the protected object inside the lock."""
        if not self._check_inner_only_in_pytest or is_pytest_env():
            self._assert_acquisition_order()
        with self._lock:
            self._owner_thread = threading.current_thread()
            try:
                with self._profiler.timeit(caller):
                    yield self._object
            finally:
                self._owner_thread = None

    def get_lock(self) -> threading.Lock:
        """Returns the underlying lock object."""
        return self._lock

    @property
    def profiler(self) -> ExecTimeProfiler:
        """Returns the profiler object."""
        return self._profiler
