"""BlockBuster module."""

from __future__ import annotations

import _thread
import asyncio
import importlib
import inspect
import io
import logging
import os
import platform
import sys
from contextlib import contextmanager
from contextvars import ContextVar
from types import ModuleType
from typing import TYPE_CHECKING, Any, TypeVar, Union


class _ModulePaths:
    def __init__(self, paths: list[str]) -> None:
        self.paths = tuple(paths)


if TYPE_CHECKING:
    import socket
    import threading
    from collections.abc import Callable, Iterable, Iterator, Sequence

    _ModuleList = Union[Sequence[Union[str, ModuleType]], _ModulePaths, None]
    _ModuleOrModuleList = Union[str, ModuleType, _ModuleList, _ModulePaths]

if platform.python_implementation() == "CPython":
    import forbiddenfruit

    HAS_FORBIDDENFRUIT = True
else:
    HAS_FORBIDDENFRUIT = False

logger = logging.getLogger(__name__)


class BlockingError(Exception):
    """BlockingError class."""

    def __init__(self, func: str) -> None:
        """Initialize BlockingError.

        Args:
            func: The blocking function.

        """
        super().__init__(f"Blocking call to {func}")


def _blocking_error(func: Callable[..., Any]) -> BlockingError:
    if inspect.isbuiltin(func):
        msg = f"Blocking call to {func.__qualname__} ({func.__self__})"
    elif inspect.ismethoddescriptor(func):
        msg = f"Blocking call to {func}"
    else:
        msg = f"Blocking call to {func.__module__}.{func.__qualname__}"
    return BlockingError(msg)


_T = TypeVar("_T")

blockbuster_skip: ContextVar[bool] = ContextVar("blockbuster_skip")


def _wrap_blocking(
    modules: Sequence[str],
    excluded_modules: Sequence[str],
    func: Callable[..., _T],
    func_name: str,
    can_block_functions: list[tuple[str, Iterable[str]]],
    can_block_predicate: Callable[..., bool],
) -> Callable[..., _T]:
    """Wrap blocking function."""

    def wrapper(*args: Any, **kwargs: Any) -> _T:
        if blockbuster_skip.get(False):
            return func(*args, **kwargs)
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return func(*args, **kwargs)
        skip_token = blockbuster_skip.set(True)
        try:
            if can_block_predicate(*args, **kwargs):
                return func(*args, **kwargs)
            frame = inspect.currentframe()
            in_test_module = False
            while frame:
                frame_filename = frame.f_code.co_filename
                if not in_test_module:
                    in_excluded_module = False
                    for excluded_module in excluded_modules:
                        if frame_filename.startswith(excluded_module):
                            in_excluded_module = True
                            break
                    if not in_excluded_module:
                        for module in modules:
                            if frame_filename.startswith(module):
                                in_test_module = True
                                break
                frame_file_name = frame_filename.replace("\\", "/")
                for filename, functions in can_block_functions:
                    if (
                        frame_file_name.endswith(filename)
                        and frame.f_code.co_name in functions
                    ):
                        return func(*args, **kwargs)
                frame = frame.f_back
            if not modules or in_test_module:
                raise BlockingError(func_name)
            return func(*args, **kwargs)
        finally:
            blockbuster_skip.reset(skip_token)

    return wrapper


def _resolve_module_paths(modules: _ModuleOrModuleList | _ModulePaths) -> _ModulePaths:
    if isinstance(modules, _ModulePaths):
        return modules
    if isinstance(modules, (str, ModuleType)):
        modules = [modules]
    resolved: list[str] = []
    for module in modules or []:
        module_ = importlib.import_module(module) if isinstance(module, str) else module
        if hasattr(module_, "__path__"):
            resolved.append(module_.__path__[0])
        elif file := module_.__file__:
            resolved.append(file)
        else:
            logger.warning("Cannot get path for %s", module_)
    return _ModulePaths(resolved)


class BlockBusterFunction:
    """BlockBusterFunction class."""

    def __init__(
        self,
        module: ModuleType | type | None,
        func_name: str,
        *,
        scanned_modules: _ModuleOrModuleList = None,
        excluded_modules: _ModuleList = None,
        can_block_functions: list[tuple[str, Iterable[str]]] | None = None,
        can_block_predicate: Callable[..., bool] = lambda *_, **__: False,
    ) -> None:
        """Create a BlockBusterFunction.

        Args:
            module: The module that contains the blocking function.
            func_name: The name of the blocking function.
            scanned_modules: The modules from which blocking calls are detected.
                If None, the blocking calls are detected from all the modules.
                Can be a module name, a module object, a list of module names or a
                list of module objects.
            excluded_modules: Sub-modules that are excluded from scanned modules.
                It doesn't mean that blocking is allowed in these modules.
                It means that they are not considered as scanned modules.
                It is helpful for instance if the tests use a pytest plugin that is
                part of the scanned modules.
                Can be a list of module names or module objects.
            can_block_functions: Optional functions in the stack where blocking is
                allowed.
            can_block_predicate: An optional predicate that determines if blocking is
                allowed.

        Raises:
            ValueError: If module is not provided and the function name doesn't contain
                a dot.

        """
        if module:
            self.module = module
            self.func_name = func_name
            self.original_func = getattr(module, func_name, None)
        else:
            tokens = func_name.split(".")
            if len(tokens) < 2:  # noqa: PLR2004
                msg = "module is required if func_name does not contain '.'"
                raise ValueError(msg)
            self.module = importlib.import_module(tokens.pop(0))
            while len(tokens) > 1:
                self.module = getattr(self.module, tokens.pop(0))
            self.func_name = tokens[0]
        self.original_func = getattr(self.module, self.func_name, None)
        if module is None:
            self.full_name = func_name
        else:
            self.full_name = f"{module.__name__}.{func_name}"
        self.can_block_functions: list[tuple[str, Iterable[str]]] = (
            can_block_functions or []
        )
        self.can_block_predicate: Callable[..., bool] = can_block_predicate
        self.activated = False
        self._scanned_modules = _resolve_module_paths(scanned_modules).paths
        self._excluded_modules = _resolve_module_paths(excluded_modules).paths

    def activate(self) -> BlockBusterFunction:
        """Activate the blocking detection.

        Returns:
            This BlockBusterFunction.
        """
        if self.original_func is None or self.activated:
            return self
        checker = _wrap_blocking(
            self._scanned_modules,
            self._excluded_modules,
            self.original_func,
            self.full_name,
            self.can_block_functions,
            self.can_block_predicate,
        )
        try:
            setattr(self.module, self.func_name, checker)
        except TypeError:
            if not HAS_FORBIDDENFRUIT:
                return self
            forbiddenfruit.curse(self.module, self.func_name, checker)
        self.activated = True
        return self

    def deactivate(self) -> BlockBusterFunction:
        """Deactivate the blocking detection.

        Returns:
            This BlockBusterFunction.
        """
        if self.original_func is None or not self.activated:
            return self
        try:
            setattr(self.module, self.func_name, self.original_func)
        except TypeError:
            if not HAS_FORBIDDENFRUIT:
                return self
            forbiddenfruit.curse(self.module, self.func_name, self.original_func)
        self.activated = False
        return self

    def can_block_in(
        self, filename: str, functions: str | Iterable[str]
    ) -> BlockBusterFunction:
        """Add functions where it is allowed to block.

        Args:
            filename (str): The filename that contains the functions.
            functions (str | Iterable[str]): The functions where blocking is allowed.

        Returns:
            This BlockBusterFunction.
        """
        if isinstance(functions, str):
            functions = {functions}
        self.can_block_functions.append((filename, functions))
        return self


def _get_time_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    return {
        "time.sleep": BlockBusterFunction(
            None,
            "time.sleep",
            can_block_functions=[
                ("/pydevd.py", {"_do_wait_suspend"}),
            ],
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
    }


def _get_os_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    functions = {
        f"os.{method}": BlockBusterFunction(
            None,
            f"os.{method}",
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
        for method in (
            "statvfs",
            "rename",
            "remove",
            "rmdir",
            "link",
            "symlink",
            "listdir",
            "access",
        )
    }

    functions["os.getcwd"] = BlockBusterFunction(
        None,
        "os.getcwd",
        can_block_functions=[
            ("coverage/control.py", {"_should_trace"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.stat"] = BlockBusterFunction(
        None,
        "os.stat",
        can_block_functions=[
            ("<frozen importlib._bootstrap>", {"_find_and_load"}),
            ("linecache.py", {"checkcache", "updatecache"}),
            ("coverage/control.py", {"_should_trace"}),
            ("coverage/python.py", {"get_python_source"}),
            ("asyncio/unix_events.py", {"create_unix_server", "_stop_serving"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.mkdir"] = BlockBusterFunction(
        None,
        "os.mkdir",
        can_block_functions=[("_pytest/assertion/rewrite.py", {"try_makedirs"})],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.replace"] = BlockBusterFunction(
        None,
        "os.replace",
        can_block_functions=[("_pytest/assertion/rewrite.py", {"_write_pyc"})],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.readlink"] = BlockBusterFunction(
        None,
        "os.readlink",
        can_block_functions=[
            ("coverage/control.py", {"_should_trace"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.sendfile"] = BlockBusterFunction(
        None,
        "os.sendfile",
        can_block_functions=[
            ("asyncio/base_events.py", {"sendfile"}),
            ("asyncio/unix_events.py", {"_sock_sendfile_native_impl"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.unlink"] = BlockBusterFunction(
        None,
        "os.unlink",
        can_block_functions=[
            ("asyncio/unix_events.py", {"_stop_serving"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    if platform.python_implementation() != "CPython" or sys.version_info >= (3, 9):
        with os.scandir() as scandir_it:
            functions["os.scandir"] = BlockBusterFunction(
                type(scandir_it),
                "__next__",
                scanned_modules=modules,
                excluded_modules=excluded_modules,
            )

    for method in (
        "ismount",
        "samestat",
        "sameopenfile",
    ):
        functions[f"os.path.{method}"] = BlockBusterFunction(
            None,
            f"os.path.{method}",
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )

    functions["os.path.islink"] = BlockBusterFunction(
        None,
        "os.path.islink",
        can_block_functions=[
            ("coverage/control.py", {"_should_trace"}),
            ("/pydevd_file_utils.py", {"get_abs_path_real_path_and_base_from_file"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    functions["os.path.abspath"] = BlockBusterFunction(
        None,
        "os.path.abspath",
        can_block_functions=[
            ("_pytest/assertion/rewrite.py", {"_should_rewrite"}),
            ("coverage/control.py", {"_should_trace"}),
            ("/pydevd_file_utils.py", {"get_abs_path_real_path_and_base_from_file"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
    )

    def os_rw_exclude(fd: int, *_: Any, **__: Any) -> bool:
        return hasattr(os, "get_blocking") and not os.get_blocking(fd)

    os_rw_kwargs = (
        {} if platform.system() == "Windows" else {"can_block_predicate": os_rw_exclude}
    )

    functions["os.read"] = BlockBusterFunction(
        None,
        "os.read",
        can_block_functions=[
            ("asyncio/base_events.py", {"subprocess_exec"}),
            ("asyncio/base_events.py", {"subprocess_shell"}),
        ],
        scanned_modules=modules,
        excluded_modules=excluded_modules,
        **os_rw_kwargs,
    )
    functions["os.write"] = BlockBusterFunction(
        None,
        "os.write",
        can_block_functions=None,
        scanned_modules=modules,
        excluded_modules=excluded_modules,
        **os_rw_kwargs,
    )

    return functions


def _get_io_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    stdout = sys.stdout
    stderr = sys.stderr

    def file_read_exclude(file: io.IOBase, *_: Any, **__: Any) -> bool:
        try:
            file.fileno()
        except io.UnsupportedOperation:
            return not file.isatty()
        return False

    def file_write_exclude(file: io.IOBase, *_: Any, **__: Any) -> bool:
        if file in {stdout, stderr, sys.stdout, sys.stderr} or file.isatty():
            return True
        try:
            file.fileno()
        except io.UnsupportedOperation:
            return True
        return False

    return {
        "io.BufferedReader.read": BlockBusterFunction(
            None,
            "io.BufferedReader.read",
            can_block_functions=[
                ("<frozen importlib._bootstrap_external>", {"get_data"}),
                ("_pytest/assertion/rewrite.py", {"_rewrite_test", "_read_pyc"}),
                ("coverage/python.py", {"read_python_source"}),
            ],
            can_block_predicate=file_read_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "io.BufferedWriter.write": BlockBusterFunction(
            None,
            "io.BufferedWriter.write",
            can_block_functions=[
                ("<frozen importlib._bootstrap>", {"_find_and_load"}),
                ("_pytest/assertion/rewrite.py", {"_write_pyc"}),
            ],
            can_block_predicate=file_write_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "io.BufferedRandom.read": BlockBusterFunction(
            None,
            "io.BufferedRandom.read",
            can_block_predicate=file_read_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "io.BufferedRandom.write": BlockBusterFunction(
            None,
            "io.BufferedRandom.write",
            can_block_predicate=file_write_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "io.TextIOWrapper.read": BlockBusterFunction(
            None,
            "io.TextIOWrapper.read",
            can_block_functions=[("aiofile/version.py", {"<module>"})],
            can_block_predicate=file_read_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "io.TextIOWrapper.write": BlockBusterFunction(
            None,
            "io.TextIOWrapper.write",
            can_block_predicate=file_write_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
    }


def _socket_exclude(sock: socket.socket, *_: Any, **__: Any) -> bool:
    return not sock.getblocking()


def _get_socket_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    return {
        f"socket.socket.{method}": BlockBusterFunction(
            None,
            f"socket.socket.{method}",
            can_block_predicate=_socket_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
        for method in (
            "connect",
            "accept",
            "send",
            "sendall",
            "sendto",
            "recv",
            "recv_into",
            "recvfrom",
            "recvfrom_into",
            "recvmsg",
        )
    }


def _get_ssl_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    return {
        f"ssl.SSLSocket.{method}": BlockBusterFunction(
            None,
            f"ssl.SSLSocket.{method}",
            can_block_predicate=_socket_exclude,
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
        for method in ("write", "send", "read", "recv")
    }


def _get_sqlite_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    functions = {
        f"sqlite3.Cursor.{method}": BlockBusterFunction(
            None,
            f"sqlite3.Cursor.{method}",
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
        for method in (
            "execute",
            "executemany",
            "executescript",
            "fetchone",
            "fetchmany",
            "fetchall",
        )
    }

    for method in ("execute", "executemany", "executescript", "commit", "rollback"):
        functions[f"sqlite3.Connection.{method}"] = BlockBusterFunction(
            None,
            f"sqlite3.Connection.{method}",
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )

    return functions


def _get_lock_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    def lock_acquire_exclude(
        lock: threading.Lock,
        blocking: bool = True,  # noqa: FBT001, FBT002
        timeout: int = -1,
    ) -> bool:
        return not blocking or timeout == 0 or not lock.locked()

    return {
        "threading.Lock.acquire": BlockBusterFunction(
            _thread.LockType,
            "acquire",
            can_block_predicate=lock_acquire_exclude,
            can_block_functions=[
                ("threading.py", {"start"}),
                ("/pydevd.py", {"_do_wait_suspend"}),
                ("asyncio/base_events.py", {"shutdown_default_executor"}),
                ("concurrent/futures/thread.py", {"submit"}),
            ],
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
        "threading.Lock.acquire_lock": BlockBusterFunction(
            _thread.LockType,
            "acquire_lock",
            can_block_predicate=lock_acquire_exclude,
            can_block_functions=[
                ("threading.py", {"start"}),
                ("concurrent/futures/thread.py", {"submit"}),
            ],
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        ),
    }


def _get_builtins_wrapped_functions(
    modules: _ModuleOrModuleList = None, excluded_modules: _ModuleList = None
) -> dict[str, BlockBusterFunction]:
    return {
        "builtins.input": BlockBusterFunction(
            None,
            "builtins.input",
            scanned_modules=modules,
            excluded_modules=excluded_modules,
        )
    }


class BlockBuster:
    """BlockBuster class."""

    def __init__(
        self,
        scanned_modules: _ModuleOrModuleList = None,
        *,
        excluded_modules: _ModuleList = None,
    ) -> None:
        """Initialize BlockBuster.

        Args:
            scanned_modules: The modules from which blocking calls are detected.
                If None, the blocking calls are detected from all the modules.
                Can be a module name, a module object, a list of module names or a
                list of module objects.
            excluded_modules: Sub-modules that are excluded from scanned modules.
                It doesn't mean that blocking is allowed in these modules.
                It means that they are not considered as scanned modules.
                It is helpful for instance if the tests use a pytest plugin that is
                part of the scanned modules.
                Can be a list of module names or module objects.
        """
        scanned_module_paths = _resolve_module_paths(scanned_modules)
        excluded_module_paths = _resolve_module_paths(excluded_modules)
        self.functions = {
            **_get_time_wrapped_functions(scanned_module_paths, excluded_module_paths),
            **_get_os_wrapped_functions(scanned_module_paths, excluded_module_paths),
            **_get_io_wrapped_functions(scanned_module_paths, excluded_module_paths),
            **_get_socket_wrapped_functions(
                scanned_module_paths, excluded_module_paths
            ),
            **_get_ssl_wrapped_functions(scanned_module_paths, excluded_module_paths),
            **_get_sqlite_wrapped_functions(
                scanned_module_paths, excluded_module_paths
            ),
            **_get_lock_wrapped_functions(scanned_module_paths, excluded_module_paths),
            **_get_builtins_wrapped_functions(
                scanned_module_paths, excluded_module_paths
            ),
        }

    def activate(self) -> None:
        """Activate all the functions."""
        activated_functions: list[BlockBusterFunction] = []
        try:
            for wrapped_function in self.functions.values():
                was_activated = wrapped_function.activated
                wrapped_function.activate()
                if wrapped_function.activated and not was_activated:
                    activated_functions.append(wrapped_function)
        except BaseException:
            for wrapped_function in reversed(activated_functions):
                wrapped_function.deactivate()
            raise

    def deactivate(self) -> None:
        """Deactivate all the functions."""
        for wrapped_function in self.functions.values():
            wrapped_function.deactivate()


@contextmanager
def blockbuster_ctx(
    scanned_modules: _ModuleOrModuleList = None, *, excluded_modules: _ModuleList = None
) -> Iterator[BlockBuster]:
    """Context manager for using BlockBuster.

    Args:
        scanned_modules: The modules from which blocking calls are detected.
            If None, the blocking calls are detected from all the modules.
            Can be a list of module names or module objects.
        excluded_modules: Sub-modules that are excluded from scanned modules.
            It doesn't mean that blocking is allowed in these modules.
            It means that they are not considered as scanned modules.
            It is helpful for instance if the tests use a pytest plugin that is
            part of the scanned modules.
            Can be a list of module names or module objects.

    Yields:
        The BlockBuster context.
    """
    blockbuster = BlockBuster(scanned_modules, excluded_modules=excluded_modules)
    blockbuster.activate()
    try:
        yield blockbuster
    finally:
        blockbuster.deactivate()
