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
