"""Misc helper functions that can be used throughout the project"""
import functools
import types
import inspect
import threading
import traceback
import pprint
from inspect import signature, Parameter
from feathersdk.utils.feathertypes import *
from enum import Enum
import time
import os
import re
from feathersdk.utils.logger import warning, error


TExc = TypeVar('TExc', bound=Exception)
TEnum = TypeVar('TEnum', bound=Enum)
T = TypeVar('T')
P = ParamSpec('P')
R = TypeVar('R')


try:
    import tqdm
    def progressbar(iterable: Iterable[T], progress: bool = True) -> Iterable[T]:
        return tqdm.tqdm(iterable, desc="Progress", disable=not progress)
except ImportError:
    def progressbar(iterable: Iterable[T], progress: bool = True) -> Iterable[T]:
        return iterable


def annotation_allows_type(annotation: Any, type_to_check: type) -> bool:
    """Check if a type annotation allows for the given type."""
    if isinstance(annotation, str):
        annotation = eval(annotation)
    
    if annotation in [Parameter.empty, Any, type_to_check]:
        return True
    
    # Recursively check if any of the union args allow the given type
    if get_type_origin(annotation) in [Union, getattr(types, 'UnionType', None)]:
        return any(annotation_allows_type(arg, type_to_check) for arg in get_type_args(annotation))
    
    return False


def make_exception_pickleable(cls: Type[TExc]) -> Type[TExc]:
    """Decorate an exception class to modify it to be pickleable.

    This is useful for exceptions that are raised in multiprocessing Pool calls as they get pickled and sent to the
    main process to be re-raised. During the unpickling process, the __init__ method is called and passed a single 
    string argument: the full exception message from the original pickled exception.This means if you create a new 
    exception class which doesn't take any arguments, or takes arguments which are used to build the exception message, 
    you'll get an error when unpickling due to incorrect number of arguments.

    This function modifies the __init__ method. When called, the first argument is expected to be `self`. For most
    calls, the class-overridden __init__ method will be called with the passed args/kwargs. However, if the new
    __init__ method is called with exactly one argument which is positional and a string, then the superclass's
    __init__ method will be called with the string argument, which is expected to be the full final exception message.

    NOTE: If there is ambiguity in the arguments to the overridden __init__ method (meaning: it is at all possible to 
    pass a single string parameter and the function call would be valid), we may not be able to tell if the argument is
    the exception message or the normal positional argument, and thus an exception will be raised. In this case, you can
    either set a non-string type annotation on what you pass in to the __init__ method, or make the argument 
    keyword-only.

    Args:
        cls: The exception class to modify. Should inherit from Exception.
    """
    old_init = cls.__init__

    # Inspect the different parameters in the __init__ method
    varpos_arg, pos_args, def_args = None, [], []
    for p in signature(old_init).parameters.values():
        if p.name == 'self':
            continue
        elif p.kind == Parameter.VAR_POSITIONAL:
            varpos_arg = p
        elif p.kind in [Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD]:
            if p.default is Parameter.empty:
                pos_args.append(p)
            else:
                def_args.append(p)
    
    # Check args for ambiguity
    if len(pos_args) <= 1:
        possible_def = [def_args[0]] if len(def_args) > 0 and len(pos_args) == 0 else []
        possible_varpos = [varpos_arg] if varpos_arg is not None and len(pos_args) == 0 else []
        params = possible_varpos + pos_args + possible_def
        if any(annotation_allows_type(p.annotation, str) for p in params):
            raise TypeError(f"Class \"{cls.__name__}\" has ambiguous function signature. Set type annotations to "
                            f"non-string type, or make the arguments keyword-only to avoid this error.")

    def new_init(self: Exception, *args: Any, **kwargs: Any) -> None:
        """Call super's init if we have exactly one arg, no kwargs, and the arg is a string.
        Otherwise call the original __init__ method with the arguments."""
        if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], str):
            return super(cls, self).__init__(*args)
        return old_init(self, *args, **kwargs)
    
    cls.__init__ = new_init
    return cls


_DEBUG_LOCK_PRINT = False
def set_debug_lock_print(debug_lock_print: bool) -> None:
    """Set whether to print debug messages when acquiring and releasing locks."""
    global _DEBUG_LOCK_PRINT
    _DEBUG_LOCK_PRINT = debug_lock_print


class LockAcquisitionError(Exception):
    """Raised when a lock cannot be acquired."""


class FeatherThreadLock:
    """A custom threading.Lock. Currently can be used to debug print messages. Call set_debug_lock_print() to enable
    
    Parameters
    ----------
    name: ``str``
        The name of the lock.
    timeout: ``float``
        The timeout in seconds to wait for the lock to be acquired.
    pre_acquire_callback: ``Optional[Callable[[], None]]``
        A callback to be called before the lock is acquired.
    post_release_callback: ``Optional[Callable[[], None]]``
        A callback to be called after the lock is released.
    """
    def __init__(self, name: str, timeout: float = 5.0, pre_acquire_callback: Optional[Callable[[], None]] = None, 
                 post_release_callback: Optional[Callable[[], None]] = None):
        self.lock = threading.Lock()
        self.name = name
        self.timeout = timeout
        self.pre_acquire_callback = pre_acquire_callback
        self.post_release_callback = post_release_callback
    
    def is_locked(self) -> bool:
        """Check if the lock is currently locked."""
        return self.lock.locked()
    
    def __enter__(self):
        if _DEBUG_LOCK_PRINT:
            self.func_name = inspect.currentframe().f_back.f_code.co_name
            print(f"Attempting to acquire lock \"{self.name}\" in function \"{self.func_name}\"", flush=True)

        if self.pre_acquire_callback is not None:
            self.pre_acquire_callback()

        self.lock.acquire(blocking=True, timeout=self.timeout)
        if not self.is_locked():
            raise LockAcquisitionError(f"Failed to acquire lock \"{self.name}\" in function \"{self.func_name}\"")

        if _DEBUG_LOCK_PRINT:
            print(f"Acquired lock \"{self.name}\" in function \"{self.func_name}\"", flush=True)
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        if self.is_locked():
            self.lock.release()
        if _DEBUG_LOCK_PRINT:
            print(f"Released lock \"{self.name}\" in function \"{self.func_name}\"", flush=True)
        
        if self.post_release_callback is not None:
            self.post_release_callback()


class CriticalException(Exception):
    """An exception that should still be raised and crash the program, even when caught for logging purposes."""
    pass


def try_log_exception_peacefully(
    func: Optional[Callable[P, R]] = None, 
    *, 
    capture_locals: Optional[Iterable[str]] = None
) -> Callable[P, Optional[R]]:
    """Decorate a function to log exceptions.
    
    If an exception is raised, it will be logged and the exception will be ignored, unless the exception is a 
    CriticalException, in which case it will be re-raised. Caught exceptions will return None.

    Can be used either like `@try_log_exception_peacefully` or `@try_log_exception_peacefully(args...)`.

    Args:
        func: Function to decorate. If not provided, will return a decorator that can be used to decorate a function.
        capture_locals: Iterable of local variable names to capture and log.
    """
    # Function was used like `@try_log_exception_peacefully(...)`
    if func is None:
        def decorator(f: Callable[P, R]) -> Callable[P, Optional[R]]:
            return try_log_exception_peacefully(f, capture_locals=capture_locals)
        return decorator

    capture_locals = list(capture_locals) if capture_locals is not None else []
    
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
        try:
            return func(*args, **kwargs)
        except Exception as e:
            cl_str = ""

            if len(capture_locals) > 0:
                if e.__traceback__ is not None:
                    frame = e.__traceback__.tb_next.tb_frame
                    cl_str = f"\n\tCaptured locals:\n"
                    local_vars = frame.f_locals

                    for local_var in capture_locals:
                        if local_var in local_vars:
                            s = pprint.pformat(local_vars[local_var], indent=8)
                            cl_str += f"\t{local_var} ({type(local_vars[local_var]).__name__}): {s}\n"
                        else:
                            cl_str += f"\t\t{local_var}: not found\n"

            exc_str = f"Exception in \"{func.__name__}\": \n{traceback.format_exc()}{cl_str}"
            warning(exc_str)

            if isinstance(e, CriticalException):
                raise e
            
            return None
    return wrapper


TimeType = float
"""The type that we use to handle time-related operations."""


def currtime() -> TimeType:
    """Get the current time in seconds.

    Currently (pun intended), this returns time.monotonic(), which returns time since computer power-on on linux.
    """
    return time.monotonic()


def timediff(start_time: TimeType) -> TimeType:
    """Get the time difference in seconds between the current time and the given start time."""
    return currtime() - start_time


def timestamped_class(cls: Type[T]) -> Type[T]:
    """Decorate a class to add a timestamp attribute to the class.
    
    The timestamp is set to the latest time that:
    - The __init__ method is called
    - The __post_init__ method is called (for dataclasses)
    - The object is created using the __new__ method (in case __init__ or __post_init__ are not called)
    """
    old_new = cls.__new__

    def new_new(cls: Type[T], *args: Any, **kwargs: Any) -> T:
        if old_new == object.__new__:
            obj = object.__new__(cls)
        else:
            obj = old_new(cls, *args, **kwargs)

        def _make_timestamped_init_func(cls: Type[T], func_name: str) -> None:
            if not hasattr(cls, func_name) or not callable(getattr(cls, func_name)):
                return
            
            old_func = getattr(cls, func_name)
            def new_func(self: cls, *args: Any, **kwargs: Any) -> None:
                old_func(self, *args, **kwargs)
                self.timestamp = currtime()
            setattr(cls, func_name, new_func)

        # Wrap the class's __init__ and __post_init__ methods, but only do it once
        if not hasattr(cls, '_timestamped_init_wrapped'):
            setattr(cls, '_timestamped_init_wrapped', True)
            _make_timestamped_init_func(cls, '__init__')
            _make_timestamped_init_func(cls, '__post_init__')
        
        # Place a manual timestamp here just in case the object never gets its __init__ or __post_init__ called
        obj.timestamp = currtime()
        
        return obj

    cls.__new__ = new_new
    return cls
    

def exponential_timeout(
    timeout_seconds: float, 
    func: Callable[P, Union[None, R]],
    func_args: P.args = (),
    func_kwargs: P.kwargs = {},
    start_delay: float = 1e-4,  # 100 microseconds
    max_delay: float = 1e-2,  # 10 milliseconds
    mult: float = 1.5, 
) -> Union[R, None]:
    """Perform an exponential timeout.

    Will run the given function, sleeping for longer and longer until the function returns a non-None value, or the
    timeout is reached (in which case a TimeoutError is raised).
    
    Args:
        timeout_seconds: The timeout in seconds.
        func: The function to execute. Should return None if the operation is not complete, and a non-None result
            if it is.
        func_args: Arguments to pass to the function.
        func_kwargs: Keyword arguments to pass to the function.
        start_delay: The starting delay in seconds.
        max_delay: The maximum delay in seconds.
        mult: Delay is multiplied by this factor after each iteration. (but not allowed to exceed max_delay)
    
    Raises:
        TimeoutError: If the function does not return a non-None value before the timeout is reached.
    """
    timeout_tt = TimeType(timeout_seconds)

    init_time = currtime()
    delay = start_delay
    while True:
        result = func(*func_args, **func_kwargs)
        if result is not None:
            return result
        time.sleep(delay)

        if timediff(init_time) > timeout_tt:
            raise TimeoutError(f"Timeout after {timediff(init_time)} seconds")
        delay = min(delay * mult, max_delay, timeout_tt - timediff(init_time))


def listdirfull(dir: str, startswith: Optional[str] = None, endswith: Optional[str] = None) -> List[str]:
    """List all files in the given directory, returning a list of full file paths.
    
    Args:
        dir: The directory to list files from.
        startswith: If provided, only list files that start with this string.
        endswith: If provided, only list files that end with this string.
    """
    return [os.path.join(dir, f) for f in os.listdir(dir) \
        if (startswith is None or f.startswith(startswith)) and (endswith is None or f.endswith(endswith))]


_TO_SNAKE_CASE_RE = re.compile(r'(?<!^)(?=[A-Z])')

def to_snake_case(name: str) -> str:
    """Convert a string to snake case."""
    return _TO_SNAKE_CASE_RE.sub(r'_', name).lower()


def dont_allow_on_fork(func):
    """Decorator to prevent a method from being called after fork."""

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if hasattr(self, 'is_forked') and self.is_forked:
            warning(
                f"Cannot call method `{self.__class__.__name__}.{func.__name__}` after fork. "
                f"{os.getpid()}, {threading.get_ident()}"
            )
            return
        return func(self, *args, **kwargs)

    return wrapper


class FeatherThread(threading.Thread):
    """A custom threading.Thread. Currently can be used to debug print messages."""
    def __init__(self, target, name=None, args=tuple(), kwargs=None, daemon=True):
        kwargs = {} if kwargs is None else kwargs
        name = name if name is not None else f"FeatherThread-{target.__name__}"
        super().__init__(target=target, name=name, args=args, kwargs=kwargs, daemon=daemon)
