from __future__ import absolute_import

import threading
import typing

from ddtrace.internal._unpatched import _threading as ddtrace_threading
from ddtrace.internal.datadog.profiling import stack_v2
from ddtrace.settings.profiling import config

from . import _lock


class _ProfiledThreadingLock(_lock._ProfiledLock):
    pass


class _ProfiledThreadingRLock(_lock._ProfiledLock):
    pass


class ThreadingLockCollector(_lock.LockCollector):
    """Record threading.Lock usage."""

    PROFILED_LOCK_CLASS = _ProfiledThreadingLock

    def _get_patch_target(self) -> typing.Type[threading.Lock]:
        return threading.Lock

    def _set_patch_target(
        self,
        value: typing.Any,
    ) -> None:
        threading.Lock = value


class ThreadingRLockCollector(_lock.LockCollector):
    """Record threading.RLock usage."""

    PROFILED_LOCK_CLASS = _ProfiledThreadingRLock

    def _get_patch_target(self) -> typing.Type[threading.RLock]:
        return threading.RLock

    def _set_patch_target(
        self,
        value: typing.Any,
    ) -> None:
        threading.RLock = value


# Also patch threading.Thread so echion can track thread lifetimes
def init_stack_v2() -> None:
    if config.stack.v2_enabled and stack_v2.is_available:
        _thread_set_native_id = ddtrace_threading.Thread._set_native_id  # type: ignore[attr-defined]
        _thread_bootstrap_inner = ddtrace_threading.Thread._bootstrap_inner  # type: ignore[attr-defined]

        def thread_set_native_id(self, *args, **kwargs):
            _thread_set_native_id(self, *args, **kwargs)
            stack_v2.register_thread(self.ident, self.native_id, self.name)

        def thread_bootstrap_inner(self, *args, **kwargs):
            _thread_bootstrap_inner(self, *args, **kwargs)
            stack_v2.unregister_thread(self.ident)

        ddtrace_threading.Thread._set_native_id = thread_set_native_id  # type: ignore[attr-defined]
        ddtrace_threading.Thread._bootstrap_inner = thread_bootstrap_inner  # type: ignore[attr-defined]

        # Instrument any living threads
        for thread_id, thread in ddtrace_threading._active.items():  # type: ignore[attr-defined]
            stack_v2.register_thread(thread_id, thread.native_id, thread.name)

        # Import _asyncio to ensure asyncio post-import wrappers are initialised
        from ddtrace.profiling import _asyncio  # noqa: F401

        _asyncio.link_existing_loop_to_current_thread()
