import copy
from .keys import DispatchKey
from .schema import OpSchema


class OperatorEntry:
    def __init__(self, name):
        self.name = name
        self.schema = None
        self.schema_obj = None
        self.error_overrides = None
        self.kernels = {}
        self.fallthrough = set()
        self.functionalize = None


class OpRegistry:
    def __init__(self):
        self._ops = {}
        self._aliases = {}

    def _canonical_name(self, name):
        if "::" in name:
            return name
        return f"aten::{name}"

    def _entry(self, name):
        name = self._canonical_name(name)
        entry = self._ops.get(name)
        if entry is None:
            entry = OperatorEntry(name)
            self._ops[name] = entry
        return entry

    def register_schema(self, name, schema):
        entry = self._entry(name)
        entry.schema = schema
        entry.schema_obj = OpSchema(schema)
        if entry.functionalize is None and name.endswith("_"):
            if any(param.mutates for param in entry.schema_obj.params):
                entry.functionalize = name[:-1]
        return entry

    def register_error_overrides(self, name, overrides):
        entry = self._entry(name)
        entry.error_overrides = overrides
        return entry

    def register_kernel(self, name, key, fn):
        entry = self._entry(name)
        if entry.schema_obj is None:
            raise RuntimeError(
                f"schema must be registered before kernel registration for op {entry.name}"
            )
        entry.kernels[key] = fn
        return entry

    def register_fallthrough(self, name, key):
        entry = self._entry(name)
        entry.fallthrough.add(key)
        return entry

    def register_alias(self, alias, target):
        self._aliases[alias] = target
        return alias

    def register_functionalize(self, name, functional_name):
        entry = self._entry(name)
        entry.functionalize = functional_name
        return entry

    def get_functionalize(self, name):
        entry = self._ops.get(self._canonical_name(name))
        if entry is None:
            return None
        return entry.functionalize

    def has(self, name):
        return self._canonical_name(name) in self._ops

    def resolve(self, name):
        return self._aliases.get(name, name)

    def register(self, name, device, fn, meta=None):
        key = resolve_dispatch_key(device)
        entry = self.register_kernel(name, key, fn)
        if meta is not None and DispatchKey.Meta not in entry.kernels:
            entry.kernels[DispatchKey.Meta] = meta
        return entry

    def get(self, name):
        return self._ops[self._canonical_name(name)]

    def snapshot(self):
        return {
            "ops": copy.deepcopy(self._ops),
            "aliases": copy.deepcopy(self._aliases),
            "global_fallthrough": copy.deepcopy(getattr(self, "_global_fallthrough", set())),
        }

    def restore(self, state):
        self._ops = copy.deepcopy(state["ops"])
        self._aliases = copy.deepcopy(state["aliases"])
        self._global_fallthrough = copy.deepcopy(state.get("global_fallthrough", set()))


registry = OpRegistry()


def _register_global_fallthroughs():
    placeholders = {
        DispatchKey.BackendSelect,
        DispatchKey.ADInplaceOrView,
        DispatchKey.AutogradOther,
        DispatchKey.AutogradCPU,
        DispatchKey.AutogradNPU,
        DispatchKey.AutogradCUDA,
        DispatchKey.AutogradXPU,
        DispatchKey.AutogradMeta,
        DispatchKey.Python,
        DispatchKey.CompositeImplicitAutograd,
        DispatchKey.CompositeExplicitAutograd,
        DispatchKey.PrivateUse1,
        DispatchKey.PrivateUse2,
        DispatchKey.PrivateUse3,
    }
    registry._global_fallthrough = placeholders


_register_global_fallthroughs()


def resolve_dispatch_key(device):
    if isinstance(device, DispatchKey):
        return device

    label = None
    if hasattr(device, "type"):
        label = str(device.type).lower()
    elif isinstance(device, str):
        label = device.lower()

    if label == "cpu":
        return DispatchKey.CPU
    if label == "npu":
        return DispatchKey.NPU
    if label == "meta":
        return DispatchKey.Meta
    if label == "cuda":
        return DispatchKey.CUDA
    if label == "mps":
        return DispatchKey.PrivateUse2

    raise ValueError(f"unsupported registration device: {device}")


_DISPATCH_KEY_STRING_MAP = {
    "CPU": DispatchKey.CPU,
    "NPU": DispatchKey.NPU,
    "CUDA": DispatchKey.CUDA,
    "PrivateUse1": DispatchKey.PrivateUse1,
    "MPS": DispatchKey.PrivateUse2,
    "PrivateUse2": DispatchKey.PrivateUse2,
    "Meta": DispatchKey.Meta,
    "Autograd": DispatchKey.Autograd,
    "AutogradCPU": DispatchKey.AutogradCPU,
    "AutogradNPU": DispatchKey.AutogradNPU,
    "AutogradCUDA": DispatchKey.AutogradCUDA,
    "AutogradPrivateUse1": DispatchKey.AutogradXPU,
    "AutogradMPS": DispatchKey.PrivateUse3,
    "AutogradPrivateUse2": DispatchKey.PrivateUse3,
    "AutogradMeta": DispatchKey.AutogradMeta,
    "AutogradOther": DispatchKey.AutogradOther,
    "AutogradXPU": DispatchKey.AutogradXPU,
    "CompositeImplicitAutograd": DispatchKey.CompositeImplicitAutograd,
    "CompositeExplicitAutograd": DispatchKey.CompositeExplicitAutograd,
    "BackendSelect": DispatchKey.BackendSelect,
    "Functionalize": DispatchKey.Functionalize,
    "ADInplaceOrView": DispatchKey.ADInplaceOrView,
}


def dispatch_key_from_string(name):
    """Map a dispatch key string to a DispatchKey enum value.

    Raises ValueError if the string is not recognized.
    """
    key = _DISPATCH_KEY_STRING_MAP.get(name)
    if key is None:
        raise ValueError(
            f"Unknown dispatch key string: {name!r}. "
            f"Valid keys: {', '.join(sorted(_DISPATCH_KEY_STRING_MAP))}"
        )
    return key
