from contextlib import contextmanager
from functools import wraps
import threading

_global_lock = threading.Lock()
_function_locks = {}
_executed_function_results = {}


class LocksByKey(object):
    """Perform synchronization by some key, for example server IP address

    In case of server IP address,
    """
    def __init__(self):
        self._locks = {}

    @contextmanager
    def lock(self, key):
        with _global_lock:
            if key not in self._locks:
                self._locks[key] = threading.Lock()
            lock = self._locks[key]
        with lock:
            yield


def synchronized(func):
    """Decorator that makes function synchronized

    That means that function can run in one thread at the same time
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        # here we use function and all its arguments as a key
        key = (func,)

        with _global_lock:
            if key not in _function_locks:
                _function_locks[key] = threading.Lock()
            lock = _function_locks[key]

        with lock:
            value = func(*args, **kwargs)
        return value

    return wrapper


def synchronized_by_args(func):
    """Decorator that makes function synchronized by arguments

    That means that function called with the same arguments can run in one thread at the same time
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        # here we use function and all its arguments as a key
        key = (func, args, frozenset(kwargs.items()))

        with _global_lock:
            if key not in _function_locks:
                _function_locks[key] = threading.Lock()
            lock = _function_locks[key]

        with lock:
            value = func(*args, **kwargs)
        return value

    return wrapper


def execute_once_synchronized_by_args(func):
    """Decorator that makes function to be executed once and synchronized for the same arguments

    The function will be executed exactly once for each set of arguments. It could be called in
    different threads. If one thread calls the function when another is already executing it,
    the first one will wait for the second one to finish, and then just "reuse" function's return value
    that we got in the second thread.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        # here we use function and all its arguments as a key
        key = (func, args, frozenset(kwargs.items()))

        with _global_lock:
            if key not in _function_locks:
                _function_locks[key] = threading.Lock()
            lock = _function_locks[key]

        with lock:
            if key in _executed_function_results:
                return _executed_function_results[key]

            value = func(*args, **kwargs)
            _executed_function_results[key] = value

        return value

    return wrapper


def synchronized_by_lock(lock):
    """Decorator that takes a threading lock for entire function execution

    :type lock: threading.Lock
    """
    def lock_synchronized(func):
        @wraps(func)
        def wrapper(*args, **kw):
            with lock:
                result = func(*args, **kw)
            return result

        return wrapper

    return lock_synchronized
