from functools import wraps
from typing import Callable

from abacus_common_logic.connectors.database import db


def method_decorator(
    decorator: Callable, name_or_names: str | list[str] | tuple[str] | set[str]
) -> Callable:
    """Apply a decorator to a specified class methods."""

    def decorate(cls):
        """Decorate the specified methods of the class."""

        names = [name_or_names] if isinstance(name_or_names, str) else name_or_names

        for name in names:
            method = getattr(cls, name)
            decorated_method = decorator(method)
            setattr(cls, name, decorated_method)

        return cls

    return decorate


def no_autoflush(func):
    """Decorator to disable autoflush for the duration of the function call."""

    @wraps(func)
    def wrapper(*args, **kwargs):
        with db.session.no_autoflush:
            return func(*args, **kwargs)

    return wrapper