# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

"""
Core utility functions that are generic and reusable across any Python application.

This module contains fundamental utility functions for common tasks including:
- Performance utilities (benchmarking, retries, timeouts)
- Decorators for common patterns
- Generic data structure operations
- String manipulation functions
- Network utilities
"""

import math
import multiprocessing
import random
import re
import signal
import socket
import string
import sys
import time
from functools import wraps
from typing import Any, TypeVar
from uuid import UUID

# Type variable for generic function signatures
T = TypeVar("T")

# Constants for boolean string conversion
TRUE_SET: set[str] = {"yes", "true", "t", "y", "1"}
FALSE_SET: set[str] = {"no", "false", "f", "n", "0"}


class TimeoutError(AssertionError):
    """
    Thrown when a timeout occurs in the `timeout` context manager.
    """

    def __init__(self, value="Timed Out"):
        self.value = value

    def __str__(self):
        return repr(self.value)


class benchmark:
    """
    Context manager for timing code execution.

    Usage:
    ```python
    with benchmark("Operation name"):
        # Code to benchmark
    ```
    """

    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, ty, val, tb):
        end = time.time()
        print(f"{self.name} : {end - self.start:.3f} seconds")
        return False


def retry(tries=5, delay=3, backoff=2):
    """
    Retry decorator with exponential back-off.

    Retries a function or method until it returns True or it exceeds the retry count.

    Args:
        tries (int): Number of times to try the function (default: 5)
        delay (int): Initial delay in seconds (default: 3)
        backoff (int): Backoff multiplier (default: 2, must be greater than 1)

    Returns:
        function: Decorated function that will be retried on failure

    Raises:
        ValueError: If backoff <= 1, tries < 0, or delay <= 0
    """
    if backoff <= 1:
        raise ValueError("back-off must be greater than 1")

    tries = math.floor(tries)
    if tries < 0:
        raise ValueError("tries must be 0 or greater")

    if delay <= 0:
        raise ValueError("delay must be greater than 0")

    def deco_retry(function):
        def f_retry(*args, **kwargs):
            mtries, mdelay = tries, delay  # make mutable

            try:
                return_value = function(*args, **kwargs)  # first attempt
            except Exception:
                return_value = False

            while mtries > 0:
                if return_value:  # Done on success
                    return True

                mtries -= 1  # consume an attempt
                time.sleep(mdelay)  # wait...
                mdelay *= backoff  # make future wait longer

                try:
                    return_value = function(*args, **kwargs)  # Try again
                except Exception:
                    return_value = False

            return False  # Ran out of tries :-(

        return f_retry  # true decorator -> decorated function

    return deco_retry  # @retry(arg[, ...]) -> true decorator


def _target(queue, function, *args, **kwargs):
    """
    Run a function with arguments and return output via a queue.

    This is a helper function for the Process created in _Timeout. It runs
    the function with positional arguments and keyword arguments and then
    returns the function's output by way of a queue. If an exception gets
    raised, it is returned to _Timeout to be raised by the value property.
    """
    try:
        queue.put((True, function(*args, **kwargs)))
    except Exception:
        queue.put((False, sys.exc_info()[1]))


class _Timeout:
    """
    Wrap a function and add a timeout (limit) attribute to it.

    Instances of this class are automatically generated by the timeout
    function. Wrapping a function allows asynchronous calls to be made
    and termination of execution after a timeout has passed.
    """

    def __init__(self, function, limit):
        """Initialize instance in preparation for being called."""
        self.__limit = limit
        self.__function = function
        self.__name__ = function.__name__
        self.__doc__ = function.__doc__
        self.__timeout = time.time()
        self.__process = multiprocessing.Process()
        self.__queue = multiprocessing.Queue()

    def __call__(self, *args, **kwargs):
        """Execute the embedded function object asynchronously."""
        self.__limit = kwargs.pop("timeout", self.__limit)
        self.__queue = multiprocessing.Queue(1)
        args = (self.__queue, self.__function) + args
        self.__process = multiprocessing.Process(target=_target, args=args, kwargs=kwargs)
        self.__process.daemon = True
        self.__process.start()
        self.__timeout = self.__limit + time.time()
        while not self.ready:
            time.sleep(0.01)
        return self.value

    def cancel(self):
        """Terminate any possible execution of the embedded function."""
        if self.__process.is_alive():
            self.__process.terminate()
        raise TimeoutError()

    @property
    def ready(self):
        """Read-only property indicating status of "value" property."""
        if self.__timeout < time.time():
            self.cancel()
        return self.__queue.full() and not self.__queue.empty()

    @property
    def value(self):
        """Read-only property containing data returned from function."""
        if self.ready is True:
            flag, load = self.__queue.get()
            if flag:
                return load
            raise load
        return None


def timeout(seconds=None, use_signals=True):
    """
    Add a timeout parameter to a function and return it.

    Args:
        seconds (float): optional time limit in seconds or fractions of a second. If None is passed, no timeout
                    is applied. This adds some flexibility to the usage: you can disable timing out depending
                    on the settings.
        use_signals (bool): flag indicating whether signals should be used for timing function out or the
                        multiprocessing. When using multiprocessing, timeout granularity is limited to
                        10ths of a second.

    Returns:
        function: The decorated function with timeout capability

    Raises:
        TimeoutError: if time limit is reached

    Usage:
    ```python
    @timeout(5)  # 5 second timeout
    def slow_function():
        ...
    ```
    """

    def decorate(function):
        if not seconds:
            return function

        if use_signals:

            def handler(_signum, _frame):
                raise TimeoutError()

            @wraps(function)
            def new_function(*args, **kwargs):
                new_seconds = kwargs.pop("timeout", seconds)
                if new_seconds:
                    old = signal.signal(signal.SIGALRM, handler)
                    signal.setitimer(signal.ITIMER_REAL, new_seconds)
                try:
                    return function(*args, **kwargs)
                finally:
                    if new_seconds:
                        signal.setitimer(signal.ITIMER_REAL, 0)
                        signal.signal(signal.SIGALRM, old)

            return new_function
        else:
            return _Timeout(function, seconds)

    return decorate


def var2bool(value: Any, raise_exc: bool = False, default: bool | None = None) -> bool | None:
    """
    Convert various input types to boolean values.

    Args:
        value: The value to convert to a boolean
        raise_exc: Whether to raise an exception for invalid values
        default: Default value to return if the value can't be converted (used when raise_exc is False)

    Returns:
        True, False, or the default value if the value can't be converted and raise_exc is False

    Raises:
        ValueError: If the value can't be converted and raise_exc is True

    Examples:
        >>> var2bool("yes")
        True
        >>> var2bool("no")
        False
        >>> var2bool("maybe")
        >>> var2bool("maybe", default=False)
        False
        >>> var2bool("maybe", default=True)
        True
        >>> var2bool("maybe", raise_exc=True)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        ValueError: Expected one of: 0, 1, f, false, n, no, t, true, y, yes
    """
    if isinstance(value, bool):
        return value
    elif isinstance(value, str):
        value = value.lower()
        if value in TRUE_SET:
            return True
        if value in FALSE_SET:
            return False

    if raise_exc:
        raise ValueError(f"Expected one of: {', '.join(sorted(TRUE_SET | FALSE_SET))}")
    return default


def strip_ansi(line: str) -> str:
    """
    Remove ANSI escape sequences from a string.

    Args:
        line: String potentially containing ANSI codes

    Returns:
        String with ANSI escape sequences removed
    """
    ansi_escape = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
    return ansi_escape.sub("", line)


def create_random_string(size=12, chars=string.ascii_letters + string.digits) -> str:
    """
    Generate a cryptographically strong random string.

    Args:
        size (int): Length of string to generate
        chars (str): Set of characters that will be used to generate the string

    Returns:
        str: Generated random string
    """
    return "".join(random.SystemRandom().choice(chars) for _ in range(size))


def remove_nones(obj: Any) -> Any:
    """
    Recursively remove None values from data structures.

    Works on lists, tuples, sets, and dictionaries.

    Args:
        obj: Input object to process

    Returns:
        Same data structure with None values removed
    """
    if isinstance(obj, list | tuple | set):
        return type(obj)(remove_nones(x) for x in obj if x is not None)
    elif isinstance(obj, dict):
        return type(obj)((remove_nones(k), remove_nones(v)) for k, v in obj.items() if k is not None and v is not None)
    else:
        return obj


def is_valid_uuid(uuid_to_test: Any, version: int = 4) -> bool:
    """
    Check if uuid_to_test is a valid UUID.

    Args:
        uuid_to_test: String to check if it's a valid UUID
        version: UUID version to check for (1, 2, 3, or 4)

    Returns:
        bool: True if uuid_to_test is a valid UUID, otherwise False

    Examples:
        >>> is_valid_uuid('c9bf9e57-1685-4c89-bafb-ff5af830be8a')
        True
        >>> is_valid_uuid('c9bf9e58')
        False
    """
    try:
        uuid_obj = UUID(str(uuid_to_test), version=version)
        # The test ensures the string representation is equal to the input
        # Normalize both to lowercase for case-insensitive comparison
        return str(uuid_obj).lower() == str(uuid_to_test).lower()
    except (ValueError, AttributeError, TypeError):
        return False


def hostname_resolves(hostname: str) -> bool:
    """
    Check if a hostname can be resolved.

    Args:
        hostname: The hostname to check

    Returns:
        bool: True if the hostname can be resolved, False otherwise
    """
    try:
        socket.gethostbyname(hostname)
        return True
    except OSError:
        return False


def chunks(li, n):
    """
    Yield successive n-sized chunks from a list.

    Args:
        li: The list to chunk
        n: Chunk size

    Yields:
        List slices of size n
    """
    for i in range(0, len(li), n):
        yield li[i : i + n]
