"""Make torch._dynamo's handler registration idempotent.

torch/_dynamo/variables/torch.py builds a local `handlers` dict and registers
functions into it from a nested closure, raising
AssertionError("Handler already registered for {fn}") when a fn turns up twice.
The closure is recreated per call, so it cannot be reached by patching module
attributes - the source file itself has to change.

The assert runs during dynamo tracing, which compressed-tensors triggers from
round_to_quantized_type_args while compressing the NVFP4 experts, and it fires
with fn == torch._C._cuda_exchangeDevice. Turning the assert into a plain
overwrite makes the registration idempotent.

install() must run before torch._dynamo is imported, so dwct/__init__.py calls
it first and nothing here imports torch (the path comes from find_spec).

The module also scans already-imported torch/compressed_tensors/transformers
modules for any module-level callable carrying the same assert and arms a
meta-path hook so later imports get scanned too. DWCT_PATCH_FORCE=0 disables it.
"""
from __future__ import annotations

import importlib
import importlib.abc
import importlib.util
import inspect
import os
import pathlib
import re
import sys

_installed = False
_MARK = "Handler already registered"
_MARKER = "dwct: idempotent registration"
_REL = "torch/_dynamo/variables/torch.py"
_ROOTS = ("torch", "compressed_tensors", "transformers", "spaces")

_BLOCK = re.compile(
    r"^(?P<ind>[ \t]*)if fn in handlers:\n"
    r"[ \t]*raise AssertionError\(.*Handler already registered.*\)\n"
    r"(?P=ind)handlers\[fn\] = handler\n",
    re.MULTILINE,
)
_LINE = re.compile(
    r"^(?P<ind>[ \t]*)raise AssertionError\(.*Handler already registered.*\)[ \t]*$",
    re.MULTILINE,
)

_wrapped = {}
_seen = set()


def _torch_file():
    spec = None
    try:
        spec = importlib.util.find_spec("torch")
    except Exception:
        spec = None
    roots = []
    if spec is not None and getattr(spec, "submodule_search_locations", None):
        roots.extend(list(spec.submodule_search_locations))
    roots.extend(sys.path)
    for root in roots:
        if not root:
            continue
        p = pathlib.Path(root) / "_dynamo" / "variables" / "torch.py"
        if p.is_file():
            return p
    return None


def _source_patch():
    p = _torch_file()
    if p is None:
        print("[patchfix] %s not found" % _REL)
        return
    try:
        src = p.read_text()
    except Exception as exc:
        print("[patchfix] read %s failed: %r" % (p, exc))
        return
    if _MARKER in src:
        print("[patchfix] source already patched: %s" % p)
        return

    m = _BLOCK.search(src)
    if m:
        repl = "%shandlers[fn] = handler  # %s\n" % (m.group("ind"), _MARKER)
        src = src[: m.start()] + repl + src[m.end():]
        how = "block"
    else:
        m = _LINE.search(src)
        if not m:
            print("[patchfix] assert text not found in %s" % p)
            return
        repl = "%spass  # %s" % (m.group("ind"), _MARKER)
        src = src[: m.start()] + repl + src[m.end():]
        how = "line"

    try:
        p.write_text(src)
    except Exception as exc:
        print("[patchfix] write %s failed: %r" % (p, exc))
        return

    for pyc in (p.parent / "__pycache__").glob("torch.cpython-*.pyc"):
        try:
            pyc.unlink()
        except Exception:
            pass
    importlib.invalidate_caches()
    print("[patchfix] source patched (%s): %s" % (how, p))
    if "torch._dynamo.variables.torch" in sys.modules:
        print("[patchfix] note: module already imported; edit lands next run")


def _tolerant(fn):
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except AssertionError as exc:
            if _MARK in str(exc):
                print("[patchfix] repeat registration ignored: %s" % exc)
                return None
            raise

    wrapper.__name__ = getattr(fn, "__name__", "wrapper")
    wrapper.__qualname__ = getattr(fn, "__qualname__", wrapper.__name__)
    wrapper._dwct_tolerant = True
    return wrapper


def _unwrap(obj):
    if isinstance(obj, (staticmethod, classmethod)):
        return obj.__func__
    return obj


def _rewrap(orig, new):
    if isinstance(orig, staticmethod):
        return staticmethod(new)
    if isinstance(orig, classmethod):
        return classmethod(new)
    return new


def _is_guard(obj):
    if not inspect.isfunction(obj):
        return False
    if id(obj) in _seen:
        return False
    _seen.add(id(obj))
    try:
        return _MARK in inspect.getsource(obj)
    except Exception:
        return False


def _register(orig, label):
    if id(orig) in _wrapped:
        return False
    _wrapped[id(orig)] = _tolerant(orig)
    where = ""
    try:
        where = " (%s:%d)" % (
            inspect.getsourcefile(orig),
            inspect.getsourcelines(orig)[1],
        )
    except Exception:
        pass
    print("[patchfix] guard %s%s" % (label, where))
    return True


def _scan_module(mod):
    name = getattr(mod, "__name__", "")
    found = 0
    try:
        items = list(vars(mod).items())
    except Exception:
        return 0
    for attr, obj in items:
        if getattr(obj, "__module__", None) == name and _is_guard(obj):
            if _register(obj, "%s.%s" % (name, attr)):
                found += 1
        elif inspect.isclass(obj) and getattr(obj, "__module__", None) == name:
            for mname, meth in list(vars(obj).items()):
                raw = _unwrap(meth)
                if _is_guard(raw):
                    if _register(raw, "%s.%s.%s" % (name, attr, mname)):
                        found += 1
    return found


def _repoint():
    n = 0
    for mod in list(sys.modules.values()):
        if mod is None:
            continue
        try:
            items = list(vars(mod).items())
        except Exception:
            continue
        for attr, obj in items:
            w = _wrapped.get(id(obj))
            if w is not None:
                try:
                    setattr(mod, attr, w)
                    n += 1
                except Exception:
                    pass
                continue
            if inspect.isclass(obj):
                for mname, meth in list(vars(obj).items()):
                    w = _wrapped.get(id(_unwrap(meth)))
                    if w is not None:
                        try:
                            setattr(obj, mname, _rewrap(meth, w))
                            n += 1
                        except Exception:
                            pass
    return n


class _Scanner(importlib.abc.MetaPathFinder):
    def find_spec(self, fullname, path=None, target=None):
        if fullname.split(".")[0] not in _ROOTS:
            return None
        spec = None
        for finder in list(sys.meta_path):
            if finder is self:
                continue
            try:
                spec = finder.find_spec(fullname, path, target)
            except Exception:
                spec = None
            if spec is not None:
                break
        if spec is None or getattr(spec, "loader", None) is None:
            return None
        loader = spec.loader
        if getattr(loader, "_dwct_scanned", False):
            return spec
        try:
            orig = loader.exec_module
        except Exception:
            return spec

        def exec_module(module, _orig=orig):
            _orig(module)
            try:
                if _scan_module(module):
                    _repoint()
            except Exception as exc:
                print(
                    "[patchfix] scan %s failed: %r"
                    % (getattr(module, "__name__", "?"), exc)
                )

        try:
            loader.exec_module = exec_module
            loader._dwct_scanned = True
        except Exception:
            pass
        return spec


def install():
    global _installed
    if _installed or os.environ.get("DWCT_PATCH_FORCE", "1") == "0":
        return
    _installed = True

    _source_patch()

    hits = 0
    for mod in list(sys.modules.values()):
        if mod is None:
            continue
        if getattr(mod, "__name__", "").split(".")[0] not in _ROOTS:
            continue
        try:
            hits += _scan_module(mod)
        except Exception:
            pass
    if hits:
        print("[patchfix] %d module-level guard(s), %d ref(s) re-pointed"
              % (hits, _repoint()))
    else:
        print("[patchfix] no module-level guard; import hook armed")

    sys.meta_path.insert(0, _Scanner())