import threading
from functools import partial

from py_zipkin import storage

_orig_Thread_start = threading.Thread.start
_orig_Thread_run = threading.Thread.run


def _Thread_pre_start(self):
    self._orig_tracer = None
    if storage.has_default_tracer():
        self._orig_tracer = storage.get_default_tracer().copy()


def _Thread_wrap_run(self, actual_run_fn):
    # This executes in the new OS thread
    if self._orig_tracer:
        # Inject our copied Tracer into our thread-local-storage
        storage.set_default_tracer(self._orig_tracer)
    try:
        actual_run_fn()
    finally:
        # I think this is probably a good idea for the same reasons the
        # parent class deletes __target, __args, and __kwargs
        if self._orig_tracer:
            del self._orig_tracer


def patch_threading():  # pragma: no cover
    """Monkey-patch threading module to work better with tracing."""

    def _new_start(self):
        _Thread_pre_start(self)
        _orig_Thread_start(self)

    def _new_run(self):
        _Thread_wrap_run(self, partial(_orig_Thread_run, self))

    threading.Thread.start = _new_start
    threading.Thread.run = _new_run


def unpatch_threading():  # pragma: no cover
    threading.Thread.start = _orig_Thread_start
    threading.Thread.run = _orig_Thread_run
