from enum import IntEnum
import threading
class DispatchKey(IntEnum):
    BackendSelect = 1 << 0
    Pipeline = 1 << 1
    Python = 1 << 2
    Functionalize = 1 << 3
    ADInplaceOrView = 1 << 4
    AutogradOther = 1 << 5
    AutogradCPU = 1 << 6
    AutogradNPU = 1 << 7
    AutogradCUDA = 1 << 8
    AutogradXPU = 1 << 9
    AutogradMeta = 1 << 10
    Autograd = 1 << 11
    Meta = 1 << 12
    NPU = 1 << 13
    CUDA = 1 << 14
    CPU = 1 << 15
    PythonDispatcher = 1 << 16
    CompositeImplicitAutograd = 1 << 17
    CompositeExplicitAutograd = 1 << 18
    Autocast = 1 << 19
    PrivateUse1 = 1 << 20
    PrivateUse2 = 1 << 21
    PrivateUse3 = 1 << 22


DISPATCH_KEY_PRIORITY = [
    DispatchKey.BackendSelect,
    DispatchKey.Pipeline,
    DispatchKey.Python,
    DispatchKey.Functionalize,
    DispatchKey.ADInplaceOrView,
    DispatchKey.AutogradOther,
    DispatchKey.AutogradCPU,
    DispatchKey.AutogradNPU,
    DispatchKey.AutogradCUDA,
    DispatchKey.AutogradXPU,
    DispatchKey.AutogradMeta,
    DispatchKey.Autograd,
    DispatchKey.Meta,
    DispatchKey.NPU,
    DispatchKey.CUDA,
    DispatchKey.CPU,
    DispatchKey.PythonDispatcher,
    DispatchKey.CompositeImplicitAutograd,
    DispatchKey.CompositeExplicitAutograd,
    DispatchKey.Autocast,
    DispatchKey.PrivateUse1,
    DispatchKey.PrivateUse2,
    DispatchKey.PrivateUse3,
]


_TLS = threading.local()


def _tls_state():
    state = getattr(_TLS, "dispatch_tls", None)
    if state is None:
        state = {"include": [], "exclude": []}
        _TLS.dispatch_tls = state
    return state


def _mask_from_keys(keys):
    if isinstance(keys, DispatchKey):
        return int(keys)
    if isinstance(keys, int):
        return int(keys)
    mask = 0
    for key in keys:
        mask |= int(key)
    return mask


def _push_mask(stack, keys):
    mask = _mask_from_keys(keys)
    stack.append(mask)
    return mask


def _pop_mask(stack):
    if stack:
        stack.pop()


class include_keys:
    def __init__(self, keys):
        self._keys = keys
        self._stack = None

    def __enter__(self):
        state = _tls_state()
        self._stack = state["include"]
        _push_mask(self._stack, self._keys)
        return self

    def __exit__(self, exc_type, exc, tb):
        _pop_mask(self._stack)


class exclude_keys:
    def __init__(self, keys):
        self._keys = keys
        self._stack = None

    def __enter__(self):
        state = _tls_state()
        self._stack = state["exclude"]
        _push_mask(self._stack, self._keys)
        return self

    def __exit__(self, exc_type, exc, tb):
        _pop_mask(self._stack)


def _effective_mask(base_mask):
    state = _tls_state()
    include_mask = 0
    for mask in state["include"]:
        include_mask |= mask
    exclude_mask = 0
    for mask in state["exclude"]:
        exclude_mask |= mask
    return (int(base_mask) | include_mask) & ~exclude_mask


def apply_tls_masks(keyset):
    if isinstance(keyset, DispatchKeySet):
        return DispatchKeySet(_effective_mask(keyset.mask))
    return DispatchKeySet(_effective_mask(keyset))


class DispatchKeySet:
    def __init__(self, mask=0):
        if isinstance(mask, (set, list, tuple)):
            # Convert set/list/tuple of DispatchKey to bitmask
            self.mask = 0
            for key in mask:
                self.mask |= int(key)
        else:
            self.mask = int(mask)

    def __contains__(self, key):
        return bool(self.mask & int(key))

    def has(self, key):
        return bool(self.mask & int(key))

    def add(self, key):
        self.mask |= int(key)
        return self

    def remove(self, key):
        self.mask &= ~int(key)
        return self

    def without(self, keys):
        mask = self.mask
        if isinstance(keys, (set, list, tuple)):
            for key in keys:
                mask &= ~int(key)
        else:
            mask &= ~int(keys)
        return DispatchKeySet(mask)

    def iter_keys(self):
        for key in DISPATCH_KEY_PRIORITY:
            if self.mask & int(key):
                yield key

    @classmethod
    def from_mask(cls, mask):
        return cls(int(mask))

    @classmethod
    def from_tensors(cls, tensors, *, grad_enabled=False, pipeline_enabled=False, functionalize_enabled=False, device=None, autocast_enabled=False):
        has_meta = False
        has_npu = False
        has_cuda = False
        has_mps = False
        has_cpu = False
        requires_grad = False
        saw_device = False
        for tensor in tensors:
            if not hasattr(tensor, "device"):
                continue
            saw_device = True
            dev = tensor.device
            dev_type = dev.type if hasattr(dev, "type") else dev
            if dev_type == "meta":
                has_meta = True
            elif dev_type == "npu":
                has_npu = True
            elif dev_type == "cuda":
                has_cuda = True
            elif dev_type == "mps":
                has_mps = True
            else:
                has_cpu = True
            if getattr(tensor, "requires_grad", False):
                requires_grad = True
        if (not saw_device) and device is not None:
            dev_type = device.type if hasattr(device, "type") else device
            if dev_type == "meta":
                has_meta = True
            elif dev_type == "npu":
                has_npu = True
            elif dev_type == "cuda":
                has_cuda = True
            elif dev_type == "mps":
                has_mps = True
            else:
                has_cpu = True
        mask = 0
        if has_meta:
            mask |= int(DispatchKey.Meta)
        elif has_npu:
            mask |= int(DispatchKey.NPU)
        elif has_cuda:
            mask |= int(DispatchKey.CUDA)
        elif has_mps:
            mask |= int(DispatchKey.PrivateUse2)
        else:
            mask |= int(DispatchKey.CPU)
        if grad_enabled and requires_grad:
            mask |= int(DispatchKey.ADInplaceOrView)
            mask |= int(DispatchKey.Autograd)
            if has_meta:
                mask |= int(DispatchKey.AutogradMeta)
            elif has_npu:
                mask |= int(DispatchKey.AutogradNPU)
            elif has_cuda:
                mask |= int(DispatchKey.AutogradCUDA)
            elif has_mps:
                mask |= int(DispatchKey.PrivateUse3)
            else:
                mask |= int(DispatchKey.AutogradCPU)
        if functionalize_enabled:
            mask |= int(DispatchKey.Functionalize)
        if autocast_enabled:
            mask |= int(DispatchKey.Autocast)
        if pipeline_enabled and not has_meta and not has_cuda:
            mask |= int(DispatchKey.Pipeline)
        return cls(mask)
