import contextlib
import copy
import hashlib
import importlib
import inspect
import io
import os
import pickle
import tokenize
import unittest
from collections.abc import Callable
from contextvars import ContextVar
from dataclasses import dataclass
from types import FunctionType, ModuleType
from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar
from typing_extensions import deprecated

from torch._utils_internal import justknobs_check


# Types saved/loaded in configs
CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict)

# Immutable scalar types that don't need deepcopy when returned from configs.
# Everything else is defensively copied to prevent accidental mutation.
_IMMUTABLE_CONFIG_TYPES = (int, float, bool, type(None), str, tuple)


# Duplicated, because mypy needs these types statically
T = TypeVar("T", bound=int | float | bool | str | list | set | tuple | dict | None)


_UNSET_SENTINEL = object()


@dataclass(kw_only=True)
class _Config(Generic[T]):
    """Represents a config with richer behaviour than just a default value.
    ::
        i.e.
        foo = Config(justknob="//foo:bar", default=False)
        install_config_module(...)

    This configs must be installed with install_config_module to be used

    Precedence Order:
        alias: If set, the directly use the value of the alias.
        env_name_force: If set, this environment variable has precedence over
            everything after this.
            If multiple env variables are given, the precedence order is from
            left to right.
        user_override: If a user sets a value (i.e. foo.bar=True), that
            has precedence over everything after this.  User overrides are thread-local.
        env_name_default: If set, this environment variable will override everything
            after this.
            If multiple env variables are given, the precedence order is from
            left to right.
        justknob: If this pytorch installation supports justknobs, that will
            override defaults, but will not override the user_override precedence.
        default: This value is the lowest precedence, and will be used if nothing is
            set.

    Environment Variables:
        These are interpreted to be either "0" or "1" to represent true and false.

    Arguments:
        justknob: the name of the feature / JK. In OSS this is unused.
        default: is the value to default this knob to in OSS.
        alias: The alias config to read instead.
        env_name_force: The environment variable, or list of, to read that is a FORCE
            environment variable. I.e. it overrides everything except for alias.
        env_name_default: The environment variable, or list of, to read that changes the
            default behaviour. I.e. user overrides take preference.
    """

    default: T | object
    justknob: str | None = None
    env_name_default: list[str] | None = None
    env_name_force: list[str] | None = None
    value_type: type | None = None
    alias: str | None = None
    # Deprecation support
    deprecated: bool = False
    deprecation_message: str | None = None

    def __post_init__(self) -> None:
        self.env_name_default = _Config.string_or_list_of_string_to_list(
            self.env_name_default
        )
        self.env_name_force = _Config.string_or_list_of_string_to_list(
            self.env_name_force
        )

        if self.alias is not None:
            if (
                self.default is not _UNSET_SENTINEL
                or self.justknob is not None
                or self.env_name_default is not None
                or self.env_name_force is not None
            ):
                raise AssertionError(
                    "if alias is set, none of {default, justknob, \
                        env_name_default and env_name_force} can be set"
                )

    @staticmethod
    def string_or_list_of_string_to_list(
        val: str | list[str] | None,
    ) -> list[str] | None:
        if val is None:
            return None
        if isinstance(val, str):
            return [val]
        if not isinstance(val, list):
            raise AssertionError(f"val is not a list, got {type(val)}")
        return val


# In runtime, we unbox the Config[T] to a T, but typechecker cannot see this,
# so in order to allow for this dynamic behavior to work correctly with
# typechecking we are going to lie to the typechecker that Config[T] returns
# a T.
if TYPE_CHECKING:

    def Config(
        default: T | object = _UNSET_SENTINEL,
        justknob: str | None = None,
        env_name_default: str | list[str] | None = None,
        env_name_force: str | list[str] | None = None,
        value_type: type | None = None,
        alias: str | None = None,
        # Deprecation support
        deprecated: bool = False,
        deprecation_message: str | None = None,
    ) -> T: ...

else:

    def Config(
        default: T | object = _UNSET_SENTINEL,
        justknob: str | None = None,
        env_name_default: str | list[str] | None = None,
        env_name_force: str | list[str] | None = None,
        value_type: type | None = None,
        alias: str | None = None,
        # Deprecation support
        deprecated: bool = False,
        deprecation_message: str | None = None,
    ) -> _Config[T]:
        return _Config(
            default=default,
            justknob=justknob,
            env_name_default=env_name_default,
            env_name_force=env_name_force,
            value_type=value_type,
            alias=alias,
            # Deprecation support
            deprecated=deprecated,
            deprecation_message=deprecation_message,
        )


def _read_env_variable(name: str) -> bool | str | None:
    value = os.environ.get(name)
    if value == "1":
        return True
    if value == "0":
        return False
    return value


def install_config_module(module: ModuleType) -> None:
    """
    Converts a module-level config into a `ConfigModule()`.

    See _config_typing.pyi for instructions on how to get the converted module to typecheck.
    """

    class ConfigModuleInstance(ConfigModule):
        # __annotations__ is written to by Sphinx autodoc
        _bypass_keys = set({"_is_dirty", "_hash_digest", "__annotations__"})

    def visit(
        source: ModuleType | type,
        dest: ModuleType | SubConfigProxy,
        prefix: str,
    ) -> None:
        """Walk the module structure and move everything to module._config"""
        type_hints = inspect.get_annotations(source)
        for key, value in list(source.__dict__.items()):
            if (
                key.startswith("__")
                or isinstance(value, (ModuleType, FunctionType))
                or (
                    hasattr(value, "__module__")
                    and (
                        value.__module__ == "typing"
                        or value.__module__.startswith("collections.abc")
                    )
                )
                # Handle from torch.utils._config_module import Config
                or (isinstance(value, type) and issubclass(value, _Config))
            ):
                continue

            name = f"{prefix}{key}"
            annotated_type = type_hints.get(key, None)
            if isinstance(value, CONFIG_TYPES):
                config[name] = _ConfigEntry(
                    _Config(default=value, value_type=annotated_type), name
                )
                if dest is module:
                    delattr(module, key)
            elif isinstance(value, _Config):
                if annotated_type is not None and value.value_type is None:
                    value.value_type = annotated_type

                config[name] = _ConfigEntry(value, name)

                if dest is module:
                    delattr(module, key)
            elif isinstance(value, type):
                if value.__module__ != module.__name__:
                    raise AssertionError(
                        f"subconfig class {value} must be defined in module {module.__name__}"
                    )
                # a subconfig with `class Blah:` syntax
                proxy = SubConfigProxy(module, f"{name}.")
                visit(value, proxy, f"{name}.")
                if dest is module:
                    setattr(dest, key, proxy)
                else:
                    dest.__dict__[key] = proxy
            else:
                raise AssertionError(f"Unhandled config {key}={value} ({type(value)})")

    config: dict[str, _ConfigEntry] = {}

    compile_ignored_keys = get_assignments_with_compile_ignored_comments(module)

    visit(module, module, "")
    module._config = config  # type: ignore[attr-defined]
    module._compile_ignored_keys = compile_ignored_keys  # type: ignore[attr-defined]
    module.__class__ = ConfigModuleInstance
    module._is_dirty = True  # type: ignore[attr-defined]
    module._hash_digest = None  # type: ignore[attr-defined]


COMPILE_IGNORED_MARKER = "@compile_ignored"


# Gets all the keys (i.e. assignments) with a @compile_ignored comment
def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str]:
    source_code = inspect.getsource(module)
    assignments = set()

    # Tokenize the source code to retrieve comments
    tokens = tokenize.tokenize(io.BytesIO(source_code.encode("utf-8")).readline)
    current_comment = "", -1
    prev_name = ""

    for token in tokens:
        if token.type == tokenize.COMMENT:
            prev_name = ""
            maybe_current = token.string.strip()
            if COMPILE_IGNORED_MARKER in maybe_current:
                if current_comment != ("", -1):
                    raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
                current_comment = maybe_current, token.start[0]
        elif token.type == tokenize.NAME:
            # Only accept the first name token, to handle if you have
            # something like foo: Bar = ...
            if not prev_name:
                prev_name = token.string
        elif token.type == tokenize.OP and token.string == "=":
            # Check if the current assignment follows a comment
            # with COMPILE_IGNORED_MARKER
            if (
                COMPILE_IGNORED_MARKER in current_comment[0]
                and current_comment[1] == token.start[0] - 1
            ):
                assignments.add(prev_name)
                current_comment = "", -1  # reset
            prev_name = ""
    if current_comment != ("", -1):
        raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
    return assignments


@dataclass
class _ConfigEntry:
    # The default value specified in the configuration
    default: Any
    # The type of the configuration value
    value_type: type
    # The value specified by the user when they overrode the configuration
    # _UNSET_SENTINEL indicates the value is not set.
    user_override: ContextVar[object]
    # The justknob to check for this config
    justknob: str | None = None
    # environment variables are read at install time
    env_value_force: Any = _UNSET_SENTINEL
    env_value_default: Any = _UNSET_SENTINEL
    # Used to work arounds bad assumptions in unittest.mock.patch
    # The code to blame is
    # https://github.com/python/cpython/blob/94a7a4e22fb8f567090514785c69e65298acca42/Lib/unittest/mock.py#L1637
    # Essentially, mock.patch requires, that if __dict__ isn't accessible
    # (which it isn't), that after delattr is called on the object, the
    # object must throw when hasattr is called. Otherwise, it doesn't call
    # setattr again.
    # Technically we'll have an intermediate state of hiding the config while
    # mock.patch is unpatching itself, but it calls setattr after the delete
    # call so the final state is correct. It's just very unintuitive.
    # upstream bug - python/cpython#126886
    hide: bool = False
    alias: str | None = None
    # Deprecation support
    deprecated: bool = False
    deprecation_message: str | None = None
    _deprecation_warned: bool = False

    def __init__(self, config: _Config, name: str) -> None:
        self.default = config.default
        self.value_type = (
            config.value_type if config.value_type is not None else type(self.default)
        )
        self.justknob = config.justknob
        self.alias = config.alias
        # Deprecation fields
        self.deprecated = config.deprecated
        self.deprecation_message = config.deprecation_message
        self._deprecation_warned = False

        self.user_override = ContextVar(name, default=_UNSET_SENTINEL)
        if config.env_name_default is not None:
            for val in config.env_name_default:
                if (env_value := _read_env_variable(val)) is not None:
                    self.env_value_default = env_value
                    break
        if config.env_name_force is not None:
            for val in config.env_name_force:
                if (env_value := _read_env_variable(val)) is not None:
                    self.env_value_force = env_value
                    break

        # Ensure justknobs and envvars are allowlisted types
        if self.justknob is not None and self.default is not None:
            if not isinstance(self.default, bool):
                raise AssertionError(
                    f"justknobs only support booleans, {self.default} is not a boolean"
                )
        if self.value_type is not None and (
            config.env_name_default is not None or config.env_name_force is not None
        ):
            if self.value_type not in (
                bool,
                str,
                Optional[bool],  # noqa: UP045
                Optional[str],  # noqa: UP045
            ):
                raise AssertionError(
                    f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
                )


class ConfigModule(ModuleType):
    # NOTE: This should be kept in sync with _config_typing.pyi.

    # The actual configuration settings.  E.g., torch._dynamo.config.debug
    # would live as "debug" in the key, and torch._inductor.config.triton.cudagraphs
    # maps as "triton.cudagraphs". See discussion on the class for meaning of various sub items
    _config: dict[str, _ConfigEntry]
    _bypass_keys: set[str]
    _compile_ignored_keys: set[str]
    _is_dirty: bool
    _hash_digest: bytes | None

    def __init__(self) -> None:
        raise NotImplementedError(
            f"use {__name__}.install_config_module(sys.modules[__name__])"
        )

    def _warn_if_deprecated(self, name: str, config: _ConfigEntry) -> None:
        """Issue deprecation warning for config if not already warned."""
        if config.deprecated and not config._deprecation_warned:
            import warnings

            msg = f"{self.__name__}.{name} is deprecated"
            if config.deprecation_message:
                msg += f" and {config.deprecation_message}"
            msg += ". It will be removed in a future version of PyTorch."
            warnings.warn(msg, FutureWarning, stacklevel=3)
            config._deprecation_warned = True

    def __setattr__(self, name: str, value: object) -> None:
        if name in self._bypass_keys:
            super().__setattr__(name, value)
        elif name not in self._config:
            raise AttributeError(f"{self.__name__}.{name} does not exist")
        else:
            # Issue deprecation warning on write (once per config)
            config = self._config[name]
            self._warn_if_deprecated(name, config)

            if config.alias is not None:
                self._set_alias_val(config, value)
            else:
                config.user_override.set(value)
                self._is_dirty = True
                config.hide = False

    def __getattr__(self, name: str) -> Any:
        try:
            config = self._config[name]

            if config.hide:
                raise AttributeError(f"{self.__name__}.{name} does not exist")

            # Issue deprecation warning on read (once per config)
            self._warn_if_deprecated(name, config)

            alias_val = self._get_alias_val(config)
            if alias_val is not _UNSET_SENTINEL:
                return alias_val

            if config.env_value_force is not _UNSET_SENTINEL:
                return config.env_value_force

            user_override = config.user_override.get()
            if user_override is not _UNSET_SENTINEL:
                return user_override

            if config.env_value_default is not _UNSET_SENTINEL:
                return config.env_value_default

            if config.justknob is not None:
                # JK only supports bools and ints
                return justknobs_check(name=config.justknob, default=config.default)

            # Reference types can still be modified, so copy them to
            # user_overrides to prevent accidental mutation of defaults.
            if not isinstance(config.default, _IMMUTABLE_CONFIG_TYPES):
                config.user_override.set(copy.deepcopy(config.default))
                return config.user_override.get()
            return config.default

        except KeyError as e:
            # make hasattr() work properly
            raise AttributeError(f"{self.__name__}.{name} does not exist") from e

    def __delattr__(self, name: str) -> None:
        self._is_dirty = True
        # must support delete because unittest.mock.patch deletes
        # then recreate things
        self._config[name].user_override.set(_UNSET_SENTINEL)
        self._config[name].hide = True

    def _get_alias_module_and_name(
        self, entry: _ConfigEntry
    ) -> tuple[ModuleType, str] | None:
        alias = entry.alias
        if alias is None:
            return None
        module_name, constant_name = alias.rsplit(".", 1)
        try:
            module = importlib.import_module(module_name)
        except ImportError as e:
            raise AttributeError(f"config alias {alias} does not exist") from e
        return module, constant_name

    def _get_alias_val(self, entry: _ConfigEntry) -> Any:
        data = self._get_alias_module_and_name(entry)
        if data is None:
            return _UNSET_SENTINEL
        module, constant_name = data
        constant_value = getattr(module, constant_name)
        return constant_value

    def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None:
        data = self._get_alias_module_and_name(entry)
        if data is None:
            raise AssertionError(
                "alias data should not be None when setting alias value"
            )
        module, constant_name = data
        setattr(module, constant_name, val)

    def _is_default(self, name: str) -> bool:
        """
        Returns true if the config is at its default value.
        configs overridden by the env are not considered default.
        """
        config_val = self._config[name]
        # The config is not overridden by the user, and the env_value_default
        # is different from the default value (meaning user has set the env to
        # change the default value).
        not_set_env_default = (
            config_val.env_value_default is _UNSET_SENTINEL
            or config_val.env_value_default == config_val.default
        )
        not_set_env_force = (
            config_val.env_value_force is _UNSET_SENTINEL
            or config_val.env_value_force == config_val.default
        )

        unset = config_val.user_override.get() is _UNSET_SENTINEL
        # Handle reference types specially to avoid spammy warnings
        if not isinstance(config_val.default, _IMMUTABLE_CONFIG_TYPES):
            unset = unset or config_val.user_override.get() == config_val.default
        return unset and not_set_env_default and not_set_env_force

    def _get_dict(
        self,
        ignored_keys: list[str] | None = None,
        ignored_prefixes: list[str] | None = None,
        skip_default: bool = False,
    ) -> dict[str, Any]:
        """Export a dictionary of current configuration keys and values.

        This function is design to provide a single point which handles
        accessing config options and exporting them into a dictionary.
        This is used by a number of different user facing export methods
        which all have slightly different semantics re: how and what to
        skip.
        If a config is aliased, it skips this config.

        Arguments:
            ignored_keys are keys that should not be exported.
            ignored_prefixes are prefixes that if a key matches should
                not be exported
            skip_default does two things. One if a key has not been modified
                it skips it.
        """
        config: dict[str, Any] = {}
        for key, entry in self._config.items():
            if entry.alias is not None:
                continue
            if ignored_keys and key in ignored_keys:
                continue
            if ignored_prefixes:
                if any(key.startswith(prefix) for prefix in ignored_prefixes):
                    continue
            if skip_default and self._is_default(key):
                continue

            # Read value directly, bypassing __getattr__ overhead
            # (deprecation warnings, alias resolution).
            user_override = entry.user_override.get()
            if entry.env_value_force is not _UNSET_SENTINEL:
                val = entry.env_value_force
            elif user_override is not _UNSET_SENTINEL:
                val = user_override
            elif entry.env_value_default is not _UNSET_SENTINEL:
                val = entry.env_value_default
            elif entry.justknob is not None:
                val = justknobs_check(name=entry.justknob, default=entry.default)
            else:
                val = entry.default
            if not isinstance(val, _IMMUTABLE_CONFIG_TYPES):
                val = copy.deepcopy(val)
            config[key] = val

        return config

    def get_type(self, config_name: str) -> type:
        return self._config[config_name].value_type

    def save_config(self) -> bytes:
        """Convert config to a pickled blob"""
        ignored_keys = getattr(self, "_save_config_ignore", [])
        return pickle.dumps(
            self._get_dict(ignored_keys=ignored_keys),
            protocol=2,
        )

    def save_config_portable(
        self, *, ignore_private_configs: bool = True
    ) -> dict[str, Any]:
        """Convert config to portable format"""
        prefixes = []
        if ignore_private_configs:
            prefixes.append("_")
        prefixes.extend(getattr(self, "_cache_config_ignore_prefix", []))
        config = self._get_dict(ignored_prefixes=prefixes)
        for key in getattr(self, "_cache_config_factory_keys", []):
            if key in config and config[key] is not None:
                instance = config[key]()
                if hasattr(instance, "uuid"):
                    config[key] = instance.uuid()
                else:
                    raise RuntimeError(
                        f"Config '{key}' is set to {config[key]} which does not "
                        f"implement uuid(). Implement uuid() for cache key "
                        f"participation."
                    )
        return config

    def codegen_config(self) -> str:
        """Convert config to Python statements that replicate current config.
        This does NOT include config settings that are at default values.
        """

        # additional imports required
        imports = set()

        def get_module_name(func: Callable, add_dot: bool) -> str:
            module_name = func.__module__
            if module_name == "builtins":
                module_name = ""
            if add_dot and module_name != "":
                module_name += "."
            return module_name

        def add_import(func: Callable) -> None:
            module_name = get_module_name(func, False)
            if module_name:
                imports.add(module_name)

        def list_of_callables_to_string(v: list | set) -> list[str]:
            return [f"{get_module_name(item, True)}{item.__name__}" for item in v]

        def importable_callable(v: Any) -> bool:
            # functools.partial has no attributes below but is a callable
            return callable(v) and hasattr(v, "__module__") and hasattr(v, "__name__")

        def get_config_line(mod, k, v) -> str:  # type: ignore[no-untyped-def]
            """
            Return a string version of the config line.
            Handle v when v is a callable, or a list/dict of callables. Add import statements for callables if necessary.
            We assume that the value of a single config won't be a mix of callables and non-callables.

            Example output:
                import logging
                import _warnings
                torch._dynamo.config.reorderable_logging_functions = { _warnings.warn, logging.warn, print }
            """
            if importable_callable(v):
                add_import(v)
                return f"{mod}.{k} = {get_module_name(v, True)}{v.__name__}"
            elif isinstance(v, (list, set)) and all(
                importable_callable(item) for item in v
            ):
                for item in v:
                    add_import(item)
                v_list = list_of_callables_to_string(v)
                if isinstance(v, list):
                    return f"{mod}.{k} = {v_list}"
                else:
                    return f"{mod}.{k} = {{ {', '.join(v_list)} }}"
            else:
                return f"{mod}.{k} = {v!r}"

        lines = []
        mod = self.__name__
        for k, v in self._get_dict(
            ignored_keys=getattr(self, "_save_config_ignore", []), skip_default=True
        ).items():
            lines.append(get_config_line(mod, k, v))
        for import_name in imports:
            lines.insert(0, f"import {import_name}")
        return "\n".join(lines)

    def get_hash(self) -> bytes:
        """Hashes the configs that are not compile_ignored"""
        if self._is_dirty or self._hash_digest is None:
            dict_to_hash = self._get_dict(ignored_keys=list(self._compile_ignored_keys))
            string_to_hash = repr(sorted(dict_to_hash.items()))
            self._hash_digest = hashlib.md5(
                string_to_hash.encode("utf-8"), usedforsecurity=False
            ).digest()
            self._is_dirty = False
        return self._hash_digest

    @deprecated(
        "`config.to_dict()` has been deprecated. It no longer changes the underlying config."
        " use `config.get_config_copy()` instead if you just want a copy of the config, or "
        "config.load_config if you need mutable access",
        category=FutureWarning,
    )
    def to_dict(self) -> dict[str, Any]:
        return self.get_config_copy()

    @deprecated(
        "`config.shallow_copy_dict()` has been deprecated. It no longer changes the underlying config."
        " use `config.get_config_copy()` instead if you just want a copy of the config, or "
        "config.load_config if you need mutable access",
        category=FutureWarning,
    )
    def shallow_copy_dict(self) -> dict[str, Any]:
        return self.get_config_copy()

    def load_config(self, maybe_pickled_config: bytes | dict[str, Any]) -> None:
        """Restore from a prior call to save_config() or shallow_copy_dict()"""
        if not isinstance(maybe_pickled_config, dict):
            config = pickle.loads(maybe_pickled_config)
        else:
            config = maybe_pickled_config
        for k, v in config.items():
            if k in self._config:
                setattr(self, k, v)
            else:
                from torch._dynamo.utils import warn_once

                warn_once(f"key {k} with value {v} is not understood by this config")

    def get_config_copy(self) -> dict[str, Any]:
        return self._get_dict()

    def get_serializable_config_copy(self) -> dict[str, Any]:
        return self._get_dict(ignored_keys=getattr(self, "_save_config_ignore", []))

    def patch(
        self,
        arg1: str | dict[str, Any] | None = None,
        arg2: Any = None,
        **kwargs: dict[str, Any],
    ) -> "ContextDecorator":
        """
        Decorator and/or context manager to make temporary changes to a config.  Note that patched settings are thread-local.

        As a decorator:

            @config.patch("name", val)
            @config.patch(name1=val1, name2=val2)
            @config.patch({"name1": val1, "name2", val2})
            def foo(...):
                ...

        As a context manager:

            with config.patch("name", val):
                ...
        """
        changes: dict[str, Any]
        if arg1 is not None:
            if arg2 is not None:
                if not isinstance(arg1, str):
                    raise AssertionError(
                        "first argument must be a string when passing 2 positional args to patch"
                    )
                # patch("key", True) syntax
                changes = {arg1: arg2}
            else:
                if not isinstance(arg1, dict):
                    raise AssertionError(
                        "first argument must be a dict when passing a single positional arg to patch"
                    )
                # patch({"key": True}) syntax
                changes = arg1
            if kwargs:
                raise AssertionError(
                    "cannot pass both positional and keyword arguments to patch"
                )
        else:
            # patch(key=True) syntax
            changes = kwargs
            if arg2 is not None:
                raise AssertionError(
                    "second positional argument is only valid when first argument is a key string"
                )
        if not isinstance(changes, dict):
            raise AssertionError(f"expected `dict` got {type(changes)}")
        config = self

        class ConfigPatch(ContextDecorator):
            def __init__(self) -> None:
                self.changes = changes
                self._prior: ContextVar[tuple[dict[str, Any], ...]] = ContextVar(
                    f"{config.__name__}.ConfigPatch[{id(self)}]",
                    default=(),
                )

            def __enter__(self) -> None:
                prior: dict[str, Any] = {}
                for key in self.changes:
                    # KeyError on invalid entry
                    prior[key] = config.__getattr__(key)
                prior_stack = self._prior.get()
                self._prior.set((*prior_stack, prior))
                try:
                    for k, v in self.changes.items():
                        config.__setattr__(k, v)
                except Exception:
                    self._prior.set(prior_stack)
                    raise

            def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore[no-untyped-def]
                prior_stack = self._prior.get()
                if not prior_stack:
                    raise AssertionError(
                        "prior should not be empty when exiting ConfigPatch"
                    )
                prior = prior_stack[-1]
                self._prior.set(prior_stack[:-1])
                for k, v in prior.items():
                    config.__setattr__(k, v)

        return ConfigPatch()

    def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any:
        """
        A lower-overhead version of patch() for things on the critical path.

        Usage:

            # do this off the critical path
            change_fn = config.make_closure_patcher(foo=True)

            ...

            revert = change_fn()
            try:
              ...
            finally:
                revert()

        """
        config = self._config

        def change() -> Callable[[], None]:
            prior = {k: config[k].user_override.get() for k in changes}
            for k, v in changes.items():
                self._config[k].user_override.set(v)

            def revert() -> None:
                for k, v in prior.items():
                    self._config[k].user_override.set(v)

            return revert

        return change


class ContextDecorator(contextlib.ContextDecorator):
    """
    Same as contextlib.ContextDecorator, but with support for
    `unittest.TestCase`
    """

    def __enter__(self) -> None:
        raise NotImplementedError("NYI")

    def __exit__(self, exc_type, exc_val, exc_tb) -> NoReturn:  # type: ignore[no-untyped-def]
        raise NotImplementedError("NYI")

    def __call__(self, func: Callable[[Any], Any]) -> Any:
        if isinstance(func, type) and issubclass(func, unittest.TestCase):

            class _TestCase(func):  # type: ignore[valid-type, misc]
                @classmethod
                def setUpClass(cls) -> None:
                    self.__enter__()
                    try:
                        super().setUpClass()
                    except Exception:
                        self.__exit__(None, None, None)
                        raise

                @classmethod
                def tearDownClass(cls) -> None:
                    try:
                        super().tearDownClass()
                    finally:
                        self.__exit__(None, None, None)

            _TestCase.__name__ = func.__name__
            _TestCase.__qualname__ = func.__qualname__
            _TestCase.__module__ = func.__module__

            return _TestCase

        return super().__call__(func)


class SubConfigProxy:
    """
    Shim to redirect to main config.
    `config.triton.cudagraphs` maps to _config["triton.cudagraphs"]
    """

    def __init__(self, config: object, prefix: str) -> None:
        # `super().__setattr__` to bypass custom `__setattr__`
        super().__setattr__("_config", config)
        super().__setattr__("_prefix", prefix)

    def __setattr__(self, name: str, value: object) -> None:
        return self._config.__setattr__(self._prefix + name, value)

    def __getattr__(self, name: str) -> Any:
        return self._config.__getattr__(self._prefix + name)

    def __delattr__(self, name: str) -> None:
        return self._config.__delattr__(self._prefix + name)


def get_tristate_env(name: str, default: Any = None) -> bool | None:
    value = os.environ.get(name)
    if value == "1":
        return True
    if value == "0":
        return False
    return default


def inherit_fields_from(parent_cls):
    def wrapper(child_cls):
        for k, v in parent_cls.__dict__.items():
            # copy fields that are not private and not overridden
            if not k.startswith("_") and k not in child_cls.__dict__:
                setattr(child_cls, k, v)
        return child_cls

    return wrapper
