"""
This module provides decorators for various purposes:

- measure_time(): Measure the execution time of a function and print it with the given print function
- handle_error(): Handle an error with different log levels
- timeout(): Raise an exception if the function runs longer than the specified timeout
- retry(): Retry a function when specific exceptions are raised, with configurable delay and max attempts
- simple_cache(): Easy cache function with parameter caching method
- abstract(): Mark a function as abstract, using LogLevels for error handling
- deprecated(): Mark a function as deprecated, using LogLevels for warning handling
- silent(): Make a function silent (disable stdout, and stderr if specified) (alternative to stouputils.ctx.Muffle)

.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/decorators_module_1.gif
  :alt: stouputils decorators examples

.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/decorators_module_2.gif
  :alt: stouputils decorators examples
"""

# Imports
import time
from collections.abc import Callable
from enum import Enum
from functools import wraps
from pickle import dumps as pickle_dumps
from traceback import format_exc
from typing import Any, Literal

from .ctx import MeasureTime, Muffle
from .print import error, progress, warning


# Execution time decorator
def measure_time(
	func: Callable[..., Any] | None = None,
	*,
	printer: Callable[..., None] = progress,
	message: str = "",
	perf_counter: bool = True,
	is_generator: bool = False
) -> Callable[..., Any]:
	""" Decorator that will measure the execution time of a function and print it with the given print function

	Args:
		func			(Callable[..., Any] | None): Function to decorate
		printer			(Callable):	Function to use to print the execution time (e.g. debug, info, warning, error, etc.)
		message			(str):		Message to display with the execution time (e.g. "Execution time of Something"),
			defaults to "Execution time of {func.__name__}"
		perf_counter	(bool):		Whether to use time.perf_counter_ns or time.time_ns
			defaults to True (use time.perf_counter_ns)
		is_generator	(bool):		Whether the function is a generator or not (default: False)
			When True, the decorator will yield from the function instead of returning it.

	Returns:
		Callable: Decorator to measure the time of the function.

	Examples:
		.. code-block:: python

			> @measure_time(printer=info)
			> def test():
			>     pass
			> test()  # [INFO HH:MM:SS] Execution time of test: 0.000ms (400ns)
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		# Set the message if not specified, else use the provided one
		new_msg: str = message if message else f"Execution time of {_get_func_name(func)}()"

		if is_generator:
			@wraps(func)
			def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
				with MeasureTime(print_func=printer, message=new_msg, perf_counter=perf_counter):
					yield from func(*args, **kwargs)
		else:
			@wraps(func)
			def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
				with MeasureTime(print_func=printer, message=new_msg, perf_counter=perf_counter):
					return func(*args, **kwargs)
		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.measure_time", func))
		return wrapper

	# Handle both @measure_time and @measure_time(printer=..., message=..., perf_counter=..., is_generator=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that handle an error with different log levels
class LogLevels(Enum):
	""" Log level for the errors in the decorator handle_error() """
	NONE = 0
	""" Do nothing """
	WARNING = 1
	""" Show as warning """
	WARNING_TRACEBACK = 2
	""" Show as warning with traceback """
	ERROR_TRACEBACK = 3
	""" Show as error with traceback """
	RAISE_EXCEPTION = 4
	""" Raise exception """

force_raise_exception: bool = False
""" If true, error_log parameter will be set to RAISE_EXCEPTION for every next handle_error calls, useful for doctests """

def handle_error(
	func: Callable[..., Any] | None = None,
	*,
	exceptions: tuple[type[BaseException], ...] | type[BaseException] = (Exception,),
	message: str = "",
	error_log: LogLevels = LogLevels.WARNING_TRACEBACK,
	sleep_time: float = 0.0
) -> Callable[..., Any]:
	""" Decorator that handle an error with different log levels.

	Args:
		func        (Callable[..., Any] | None):    	Function to decorate
		exceptions	(tuple[type[BaseException]], ...):	Exceptions to handle
		message		(str):								Message to display with the error. (e.g. "Error during something")
		error_log	(LogLevels):						Log level for the errors
			LogLevels.NONE:					None
			LogLevels.WARNING:				Show as warning
			LogLevels.WARNING_TRACEBACK:	Show as warning with traceback
			LogLevels.ERROR_TRACEBACK:		Show as error with traceback
			LogLevels.RAISE_EXCEPTION:		Raise exception
		sleep_time	(float):							Time to sleep after the error (e.g. 0.0 to not sleep, 1.0 to sleep for 1 second)

	Examples:
		>>> @handle_error
		... def might_fail():
		...     raise ValueError("Let's fail")

		>>> @handle_error(error_log=LogLevels.WARNING)
		... def test():
		...     raise ValueError("Let's fail")
		>>> # test()	# [WARNING HH:MM:SS] Error during test: (ValueError) Let's fail
	"""
	# Update error_log if needed
	if force_raise_exception:
		error_log = LogLevels.RAISE_EXCEPTION

	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		if message != "":
			msg: str = f"{message}, "
		else:
			msg: str = message

		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			try:
				return func(*args, **kwargs)
			except exceptions as e:
				if error_log == LogLevels.WARNING:
					warning(f"{msg}Error during {_get_func_name(func)}(): ({type(e).__name__}) {e}")
				elif error_log == LogLevels.WARNING_TRACEBACK:
					warning(f"{msg}Error during {_get_func_name(func)}():\n{format_exc()}")
				elif error_log == LogLevels.ERROR_TRACEBACK:
					error(f"{msg}Error during {_get_func_name(func)}():\n{format_exc()}", exit=True)
				elif error_log == LogLevels.RAISE_EXCEPTION:
					raise e

				# Sleep for the specified time, only if the error_log is not ERROR_TRACEBACK (because it's blocking)
				if sleep_time > 0.0 and error_log != LogLevels.ERROR_TRACEBACK:
					time.sleep(sleep_time)
		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.handle_error", func))
		return wrapper

	# Handle both @handle_error and @handle_error(exceptions=..., message=..., error_log=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that raises an exception if the function runs too long
def timeout(
	func: Callable[..., Any] | None = None,
	*,
	seconds: float = 60.0,
	message: str = ""
) -> Callable[..., Any]:
	""" Decorator that raises a TimeoutError if the function runs longer than the specified timeout.

	Note: This decorator uses SIGALRM on Unix systems, which only works in the main thread.
	On Windows or in non-main threads, it will fall back to a polling-based approach.

	Args:
		func		(Callable[..., Any] | None):	Function to apply timeout to
		seconds		(float):						Timeout duration in seconds (default: 60.0)
		message		(str):							Custom timeout message (default: "Function '{func_name}' timed out after {seconds} seconds")

	Returns:
		Callable[..., Any]: Decorator that enforces timeout on the function

	Raises:
		TimeoutError: If the function execution exceeds the timeout duration

	Examples:
		>>> @timeout(seconds=2.0)
		... def slow_function():
		...     time.sleep(5)
		>>> slow_function()  # Raises TimeoutError after 2 seconds
		Traceback (most recent call last):
			...
		TimeoutError: Function 'slow_function()' timed out after 2.0 seconds

		>>> @timeout(seconds=1.0, message="Custom timeout message")
		... def another_slow_function():
		...     time.sleep(3)
		>>> another_slow_function()  # Raises TimeoutError after 1 second
		Traceback (most recent call last):
			...
		TimeoutError: Custom timeout message
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		# Check if we can use signal-based timeout (Unix only)
		import os
		use_signal: bool = os.name != "nt"  # Not Windows

		if use_signal:
			try:
				import signal
				# Verify SIGALRM is available
				use_signal = hasattr(signal, 'SIGALRM')
			except ImportError:
				use_signal = False

		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			# Build timeout message
			msg: str = message if message else f"Function '{_get_func_name(func)}()' timed out after {seconds} seconds"

			# Use signal-based timeout on Unix (main thread only)
			if use_signal:
				import signal
				import threading

				# Signal only works in main thread
				if threading.current_thread() is threading.main_thread():
					def timeout_handler(signum: int, frame: Any) -> None:
						raise TimeoutError(msg)

					# Set the signal handler and alarm
					old_handler = signal.signal(signal.SIGALRM, timeout_handler) # type: ignore
					signal.setitimer(signal.ITIMER_REAL, seconds) # type: ignore

					try:
						result = func(*args, **kwargs)
					finally:
						# Cancel the alarm and restore the old handler
						signal.setitimer(signal.ITIMER_REAL, 0) # type: ignore
						signal.signal(signal.SIGALRM, old_handler) # type: ignore

					return result

			# Fall back to polling-based timeout (Windows or non-main thread)
			import threading

			result_container: list[Any] = []
			exception_container: list[BaseException] = []

			def target() -> None:
				try:
					result_container.append(func(*args, **kwargs))
				except BaseException as e_2:
					exception_container.append(e_2)

			thread = threading.Thread(target=target, daemon=True)
			thread.start()
			thread.join(timeout=seconds)

			if thread.is_alive():
				# Thread is still running, timeout occurred
				raise TimeoutError(msg)

			# Check if an exception was raised in the thread
			if exception_container:
				raise exception_container[0]

			# Return the result if available
			if result_container:
				return result_container[0]

		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.timeout", func))
		return wrapper

	# Handle both @timeout and @timeout(seconds=..., message=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that retries a function when specific exceptions are raised
def retry(
	func: Callable[..., Any] | None = None,
	*,
	exceptions: tuple[type[BaseException], ...] | type[BaseException] = (Exception,),
	max_attempts: int = 10,
	delay: float = 1.0,
	backoff: float = 1.0,
	message: str = ""
) -> Callable[..., Any]:
	""" Decorator that retries a function when specific exceptions are raised.

	Args:
		func			(Callable[..., Any] | None):		Function to retry
		exceptions		(tuple[type[BaseException], ...]):	Exceptions to catch and retry on
		max_attempts	(int | None):						Maximum number of attempts (None for infinite retries)
		delay			(float):							Initial delay in seconds between retries (default: 1.0)
		backoff			(float):							Multiplier for delay after each retry (default: 1.0 for constant delay)
		message			(str):								Custom message to display before ", retrying" (default: "{ExceptionName} encountered while running {func_name}")

	Returns:
		Callable[..., Any]: Decorator that retries the function on specified exceptions

	Examples:
		>>> import os
		>>> @retry(exceptions=PermissionError, max_attempts=3, delay=0.1)
		... def write_file():
		...     with open("test.txt", "w") as f:
		...         f.write("test")

		>>> @retry(exceptions=(OSError, IOError), delay=0.5, backoff=2.0)
		... def network_call():
		...     pass

		>>> @retry(max_attempts=5, delay=1.0)
		... def might_fail():
		...     pass
	"""
	# Normalize exceptions to tuple
	if not isinstance(exceptions, tuple):
		exceptions = (exceptions,)

	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			attempt: int = 0
			current_delay: float = delay

			while True:
				attempt += 1
				try:
					return func(*args, **kwargs)
				except exceptions as e:
					# Check if we should retry or give up
					if max_attempts != 1 and attempt >= max_attempts:
						raise e

					# Log retry attempt
					if message:
						warning(f"{message}, retrying ({attempt + 1}/{max_attempts}): {e}")
					else:
						warning(f"{type(e).__name__} encountered while running {_get_func_name(func)}(), retrying ({attempt + 1}/{max_attempts}): {e}")

					# Wait before next attempt
					time.sleep(current_delay)
					current_delay *= backoff

		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.retry", func))
		return wrapper

	# Handle both @retry and @retry(exceptions=..., max_attempts=..., delay=...)
	if func is None:
		return decorator
	return decorator(func)

# Easy cache function with parameter caching method
def simple_cache(
	func: Callable[..., Any] | None = None,
	*,
	method: Literal["str", "pickle"] = "str"
) -> Callable[..., Any]:
	""" Decorator that caches the result of a function based on its arguments.

	The str method is often faster than the pickle method (by a little) but not as accurate with complex objects.

	Args:
		func   (Callable[..., Any] | None): Function to cache
		method (Literal["str", "pickle"]):  The method to use for caching.
	Returns:
		Callable[..., Any]: A decorator that caches the result of a function.
	Examples:
		>>> @simple_cache
		... def test1(a: int, b: int) -> int:
		...     return a + b

		>>> @simple_cache(method="str")
		... def test2(a: int, b: int) -> int:
		...     return a + b
		>>> test2(1, 2)
		3
		>>> test2(1, 2)
		3
		>>> test2(3, 4)
		7

		>>> @simple_cache
		... def factorial(n: int) -> int:
		...     return n * factorial(n - 1) if n else 1
		>>> factorial(10)   # no previously cached result, makes 11 recursive calls
		3628800
		>>> factorial(5)    # no new calls, just returns the cached result
		120
		>>> factorial(12)   # two new recursive calls, factorial(10) is cached
		479001600
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		# Create the cache dict
		cache_dict: dict[Any, Any] = {}

		# Create the wrapper
		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:

			# Get the hashed key
			if method == "str":
				hashed = str(args) + str(kwargs)
			elif method == "pickle":
				hashed = pickle_dumps((args, kwargs))
			else:
				raise ValueError("Invalid caching method. Supported methods are 'str' and 'pickle'.")

			# If the key is in the cache, return it
			if hashed in cache_dict:
				return cache_dict[hashed]

			# Else, call the function and add the result to the cache
			else:
				result: Any = func(*args, **kwargs)
				cache_dict[hashed] = result
				return result

		# Return the wrapper
		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.simple_cache", func))
		return wrapper

	# Handle both @simple_cache and @simple_cache(method=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that marks a function as abstract
def abstract(
	func: Callable[..., Any] | None = None,
	*,
	error_log: LogLevels = LogLevels.RAISE_EXCEPTION
) -> Callable[..., Any]:
	""" Decorator that marks a function as abstract.

	Contrary to the abstractmethod decorator from the abc module that raises a TypeError
	when you try to instantiate a class that has abstract methods, this decorator raises
	a NotImplementedError ONLY when the decorated function is called, indicating that the function
	must be implemented by a subclass.

	Args:
		func                (Callable[..., Any] | None): The function to mark as abstract
		error_log           (LogLevels):                 Log level for the error handling
			LogLevels.NONE:              None
			LogLevels.WARNING:           Show as warning
			LogLevels.WARNING_TRACEBACK: Show as warning with traceback
			LogLevels.ERROR_TRACEBACK:   Show as error with traceback
			LogLevels.RAISE_EXCEPTION:   Raise exception

	Returns:
		Callable[..., Any]: Decorator that raises NotImplementedError when called

	Examples:
		>>> class Base:
		...     @abstract
		...     def method(self):
		...         pass
		>>> Base().method()
		Traceback (most recent call last):
			...
		NotImplementedError: Function 'method()' is abstract and must be implemented by a subclass
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		message: str = f"Function '{_get_func_name(func)}()' is abstract and must be implemented by a subclass"
		if not func.__doc__:
			func.__doc__ = message

		@wraps(func)
		@handle_error(exceptions=NotImplementedError, error_log=error_log)
		def not_implemented_error(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			raise NotImplementedError(message)
		_set_wrapper_name(not_implemented_error, _get_wrapper_name("stouputils.decorators.abstract", func))
		return not_implemented_error

	# Handle both @abstract and @abstract(error_log=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that marks a function as deprecated
def deprecated(
	func: Callable[..., Any] | None = None,
	*,
	message: str = "",
	version: str = "",
	error_log: LogLevels = LogLevels.WARNING
) -> Callable[..., Any]:
	""" Decorator that marks a function as deprecated.

	Args:
		func        (Callable[..., Any] | None): Function to mark as deprecated
		message     (str):                       Additional message to display with the deprecation warning
		version     (str):                       Version since when the function is deprecated (e.g. "v1.2.0")
		error_log   (LogLevels):                 Log level for the deprecation warning
			LogLevels.NONE:              None
			LogLevels.WARNING:           Show as warning
			LogLevels.WARNING_TRACEBACK: Show as warning with traceback
			LogLevels.ERROR_TRACEBACK:   Show as error with traceback
			LogLevels.RAISE_EXCEPTION:   Raise exception
	Returns:
		Callable[..., Any]: Decorator that marks a function as deprecated

	Examples:
		>>> @deprecated
		... def old_function():
		...     pass

		>>> @deprecated(message="Use 'new_function()' instead", error_log=LogLevels.WARNING)
		... def another_old_function():
		...     pass
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			# Build deprecation message
			msg: str = f"Function '{_get_func_name(func)}()' is deprecated"
			if version:
				msg += f" since {version}"
			if message:
				msg += f". {message}"

			# Handle deprecation warning based on log level
			if error_log == LogLevels.WARNING:
				warning(msg)
			elif error_log == LogLevels.WARNING_TRACEBACK:
				warning(f"{msg}\n{format_exc()}")
			elif error_log == LogLevels.ERROR_TRACEBACK:
				error(f"{msg}\n{format_exc()}", exit=True)
			elif error_log == LogLevels.RAISE_EXCEPTION:
				raise DeprecationWarning(msg)

			# Call the original function
			return func(*args, **kwargs)
		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.deprecated", func))
		return wrapper

	# Handle both @deprecated and @deprecated(message=..., error_log=...)
	if func is None:
		return decorator
	return decorator(func)

# Decorator that make a function silent (disable stdout)
def silent(
	func: Callable[..., Any] | None = None,
	*,
	mute_stderr: bool = False
) -> Callable[..., Any]:
	""" Decorator that makes a function silent (disable stdout, and stderr if specified).

	Alternative to stouputils.ctx.Muffle.

	Args:
		func			(Callable[..., Any] | None):	Function to make silent
		mute_stderr		(bool):							Whether to mute stderr or not

	Examples:
		>>> @silent
		... def test():
		...     print("Hello, world!")
		>>> test()

		>>> @silent(mute_stderr=True)
		... def test2():
		...     print("Hello, world!")
		>>> test2()

		>>> silent(print)("Hello, world!")
	"""
	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
		@wraps(func)
		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
			# Use Muffle context manager to silence output
			with Muffle(mute_stderr=mute_stderr):
				return func(*args, **kwargs)
		_set_wrapper_name(wrapper, _get_wrapper_name("stouputils.decorators.silent", func))
		return wrapper

	# Handle both @silent and @silent(mute_stderr=...)
	if func is None:
		return decorator
	return decorator(func)



# "Private" functions
def _get_func_name(func: Callable[..., Any]) -> str:
	""" Get the name of a function, returns "<unknown>" if the name cannot be retrieved. """
	try:
		return func.__name__
	except AttributeError:
		return "<unknown>"

def _get_wrapper_name(decorator_name: str, func: Callable[..., Any]) -> str:
	""" Get a descriptive name for a wrapper function.

	Args:
		decorator_name	(str):					Name of the decorator
		func			(Callable[..., Any]):	Function being decorated
	Returns:
		str: Combined name for the wrapper function (e.g., "stouputils.decorators.handle_error@function_name")
	"""
	func_name: str = _get_func_name(func)

	# Remove "stouputils.decorators.*" prefix if present
	if func_name.startswith("stouputils.decorators."):
		func_name = func_name.split(".", 2)[-1]

	return f"{decorator_name}@{func_name}"


def _set_wrapper_name(wrapper: Callable[..., Any], name: str) -> None:
	""" Set the wrapper function's visible name, qualname and code object name for clearer tracebacks.

	Args:
		wrapper	(Callable[..., Any]):	Wrapper function to update
		name	(str):					New name to set
	"""
	# __name__ affects repr and some introspection
	wrapper.__name__ = name

	# __qualname__ helps when using nested scopes and in some introspection tools
	try:
		wrapper.__qualname__ = name
	except Exception:
		pass

	# Update the code object's co_name so tracebacks show the new name
	try:
		wrapper.__code__ = wrapper.__code__.replace(co_name=name)
	except Exception:
		# If code.replace isn't available, ignore silently
		pass

