"""Tagging what we install, so it can be removed again.

ComfyUI offers two shapes of hook. Some are keyed -- `add_wrapper_with_key` takes
a name and removing it is exact. Others are plain lists on `model_options`, where
`set_model_sampler_pre_cfg_function` appends an anonymous callable and nothing
records who put it there.

For the second kind, install-then-forget means a model run through the loader
twice carries two copies and applies the effect at double strength, reporting
once. That is the accumulation this project has already paid for: v4 registered
hooks on shared blocks and only a restart cleared them.

So everything we install is tagged, and everything tagged is stripped before the
next install. The walk below is generic -- it looks for our tag rather than for a
list of known hook names, because a list of hook names is a thing to keep current
and the day it falls behind is the day a hook stops being removable.
"""

from typing import Any

TAG = "_funpack_key"


def tag(fn, key: str):
    """Mark a callable as ours, so `strip` can find it later."""
    try:
        setattr(fn, TAG, key)
    except AttributeError:
        # A builtin or a slotted object cannot carry the mark, and something
        # unremovable must not be installed silently.
        raise TypeError(
            f"{fn!r} cannot be tagged, so it could never be removed again. "
            f"Install a plain function or a closure.")
    return fn


def _ours(value: Any, prefix: str) -> bool:
    key = getattr(value, TAG, None)
    return isinstance(key, str) and key.startswith(prefix)


def strip(patcher, prefix: str) -> int:
    """Remove every tagged hook under `prefix`, and nothing else.

    Covers both shapes: keyed wrappers and callbacks, and the anonymous function
    lists on model_options. Returns how many were removed, so a caller can say
    that it replaced something rather than leaving it to be inferred.
    """
    removed = 0

    for holder in (getattr(patcher, "wrappers", None), getattr(patcher, "callbacks", None)):
        if not isinstance(holder, dict):
            continue
        for by_key in holder.values():
            if not isinstance(by_key, dict):
                continue
            for key in [k for k in by_key if isinstance(k, str) and k.startswith(prefix)]:
                by_key.pop(key, None)
                removed += 1

    options = getattr(patcher, "model_options", None)
    if isinstance(options, dict):
        removed += _strip_options(options, prefix)

    # add_object_patch entries (a module attribute swapped at load time) are
    # keyed by attribute path, so a clone carries ours forward too; without this
    # a feature switched OFF kept its patch on a model that went round again.
    objects = getattr(patcher, "object_patches", None)
    if isinstance(objects, dict):
        for name in [n for n, v in objects.items() if _ours(v, prefix)]:
            objects.pop(name, None)
            removed += 1
    return removed


def _strip_options(options: dict, prefix: str) -> int:
    """Walk model_options for tagged callables, in lists or on their own."""
    removed = 0
    for name, value in list(options.items()):
        if isinstance(value, list):
            keep = [item for item in value if not _ours(item, prefix)]
            if len(keep) != len(value):
                removed += len(value) - len(keep)
                options[name] = keep
        elif callable(value) and _ours(value, prefix):
            # A chained hook remembers what it wrapped (dit_hooks.PREV): put that
            # back rather than dropping it, or stripping ours deletes someone
            # else's hook along with it.
            while callable(value) and _ours(value, prefix):
                value = getattr(value, "_funpack_prev", None)
                removed += 1
            if value is None:
                options.pop(name, None)
            else:
                options[name] = value
        elif isinstance(value, dict):
            removed += _strip_options(value, prefix)
        elif _ours(value, prefix):
            # A tagged object that is not a callable (a handler core looks up by key): it has
            # nothing underneath to put back.
            options.pop(name, None)
            removed += 1
    return removed


# --- keeping a module from killing a run ----------------------------------
#
# Catching a module that fails to INSTALL is easy and was already done. The one
# that matters is a module that installs fine and raises on step 3 of 30: the
# generation dies after the GPU time is spent, and on a rental that is the whole
# cost of the run. A modifier is an opinion about the picture. It is never worth
# the picture.
#
# So a failing hook is dropped for the rest of the run and the run continues
# without it. Loudly -- a dropped modifier is the difference between "I asked for
# this" and "this happened", so it is said once, with the traceback, and carried
# in the run's status.

#: Set by core/control.py: told when a module's hook raised, so the fault outlives the run.
on_fault = None


class Dropped:
    """What was disabled during a run, and why. One per run, not per process."""

    def __init__(self):
        self.reasons: dict = {}

    def clear(self) -> None:
        """Forget what was dropped. Called when a generation starts: this run's record starts
        empty (a module that FAILED also stays off across runs: that is
        core/control.py's quarantine, not this record) -- and because ComfyUI's cache hands the same
        object to both."""
        self.reasons.clear()

    def record(self, key: str, exc: BaseException) -> bool:
        """True the first time this key fails, so callers report once."""
        if key in self.reasons:
            return False
        self.reasons[key] = f"{type(exc).__name__}: {exc}"
        if on_fault is not None:
            try:
                on_fault(key, exc)
            except Exception:                    # noqa: BLE001 -- bookkeeping must not end a run
                pass
        return True

    def __contains__(self, key) -> bool:
        return key in self.reasons

    def __bool__(self) -> bool:
        return bool(self.reasons)

    def items(self):
        return self.reasons.items()


def guard(fn, key: str, neutral, dropped: Dropped):
    """`fn`, but unable to end a run.

    `neutral(*args, **kwargs)` produces the result the caller would have got had
    the hook never been installed -- it differs per hook shape, which is why it
    is supplied rather than guessed. After the first failure the hook is skipped
    outright: a modifier that raises once will raise every step, and thirty
    tracebacks say nothing the first one did not.
    """
    import traceback

    def guarded(*args, **kwargs):
        if key in dropped:
            return neutral(*args, **kwargs)
        try:
            return fn(*args, **kwargs)
        except Exception as exc:                 # noqa: BLE001 -- the whole point
            if dropped.record(key, exc):
                from . import log
                log.warning(
                    key,
                    "failed during sampling and is now OFF for the rest of this run; "
                    "the run continues without it\n"
                    + "".join(traceback.format_exception(exc)).rstrip())
            return neutral(*args, **kwargs)

    # Carried over so a guarded hook is still ours to strip.
    marker = getattr(fn, TAG, None)
    if marker is not None:
        guarded.__dict__[TAG] = marker
    return guarded


class _Executor:
    """ComfyUI's wrapper executor, seen through a glass: calls and attributes pass
    through, and what the call raised or returned is remembered."""

    def __init__(self, executor):
        self._executor = executor
        self.raised = None
        self.done = False
        self.result = None

    def __call__(self, *args, **kwargs):
        try:
            self.result = self._executor(*args, **kwargs)
        except BaseException as exc:
            self.raised = exc
            raise
        self.done = True
        return self.result

    def __getattr__(self, name):
        return getattr(self._executor, name)


def guard_wrapper(fn, key: str, dropped: Dropped):
    """`guard` for `wrapper(executor, *args)` hooks, which the generic one gets
    wrong twice: the executor IS the rest of the run, so

    * what the executor raises (an interrupt, an OOM, the next module's bug) is
      not this module's failure -- it passes through untouched, and
    * once the executor has run, falling back to it would run the whole thing
      AGAIN -- a wrapper that fails afterwards returns what already ran.
    """
    import traceback

    def guarded(executor, *args, **kwargs):
        if key in dropped:
            return executor(*args, **kwargs)
        seen = _Executor(executor)
        try:
            return fn(seen, *args, **kwargs)
        except Exception as exc:                 # noqa: BLE001
            if exc is seen.raised:
                raise
            if dropped.record(key, exc):
                from . import log
                log.warning(
                    key,
                    "failed during sampling and is now OFF for the rest of this run; "
                    "the run continues without it\n"
                    + "".join(traceback.format_exception(exc)).rstrip())
            return seen.result if seen.done else executor(*args, **kwargs)

    marker = getattr(fn, TAG, None)
    if marker is not None:
        guarded.__dict__[TAG] = marker
    return guarded


def guard_unet(fn, key: str, dropped: Dropped, prior=None):
    """`guard` for the `wrapper(apply_fn, args)` hook. `apply_fn` is the model: what it raises (an OOM,
    an interrupt) is not this module's failure and passes through; when the module itself fails it is
    dropped and the call goes to whatever wrapper was installed before it (`prior`), else the model."""
    import traceback

    def fallback(apply_fn, args):
        if prior is not None:
            return prior(apply_fn, args)
        return apply_fn(args["input"], args["timestep"], **args.get("c", {}))

    def guarded(apply_fn, args):
        if key in dropped:
            return fallback(apply_fn, args)
        seen = _Executor(apply_fn)
        try:
            return fn(seen, args)
        except Exception as exc:                 # noqa: BLE001
            if exc is seen.raised:
                raise
            if dropped.record(key, exc):
                from . import log
                log.warning(
                    key,
                    "failed during sampling and is now OFF for the rest of this run; "
                    "the run continues without it\n"
                    + "".join(traceback.format_exception(exc)).rstrip())
            return seen.result if seen.done else fallback(apply_fn, args)

    marker = getattr(fn, TAG, None)
    if marker is not None:
        guarded.__dict__[TAG] = marker
    return guarded


# The neutral result for each hook ComfyUI offers, by the method that installs
# it. Core knows ComfyUI's own shapes here -- not FunPack's features -- because
# "what this hook returns when it does nothing" is a fact about ComfyUI.
def _neutral_wrapper(executor, *args, **kwargs):
    return executor(*args, **kwargs)


NEUTRAL = {
    "add_wrapper_with_key": _neutral_wrapper,
    "add_callback_with_key": lambda *a, **k: None,
    "set_model_sampler_pre_cfg_function": lambda args: args["conds_out"],
    "set_model_sampler_post_cfg_function": lambda args: args["denoised"],
    "set_model_sampler_cfg_function": lambda args: args["input"] - (
        args["uncond_denoised"] + (args["cond_denoised"] - args["uncond_denoised"]) * args["cond_scale"]),
    "set_model_unet_function_wrapper": lambda apply_fn, args: apply_fn(args["input"], args["timestep"],
                                                                       **args.get("c", {})),
    "set_model_denoise_mask_function": lambda sigma, denoise_mask, extra_options=None: denoise_mask,
}

# Which positional argument holds the callable, for methods that take more than one.
CALLABLE_AT = {"add_wrapper_with_key": 2, "add_callback_with_key": 2}


class GuardedPatcher:
    """A ModelPatcher that guards every hook installed through it.

    A module is handed one of these instead of the real patcher, so it does not
    opt in to being safe -- a module that could forget to is a module that can
    break the foundation. Everything else forwards untouched.
    """

    def __init__(self, patcher, key: str, dropped: Dropped, guarding: bool = True):
        object.__setattr__(self, "_patcher", patcher)
        object.__setattr__(self, "_key", key)
        object.__setattr__(self, "_dropped", dropped)
        object.__setattr__(self, "_guarding", guarding)
        object.__setattr__(self, "unguarded", [])

    def __getattr__(self, name):
        target = getattr(self._patcher, name)
        if not callable(target) or not name.startswith(("set_model_", "add_wrapper", "add_callback")):
            return target

        if not self._guarding:
            # Full control: install exactly what the module wrote, and let it
            # fail where it fails. Someone asking for the raw behaviour wants
            # the real traceback at the real step, not a tidy recovery.
            return target

        neutral = NEUTRAL.get(name)
        if neutral is None:
            # Honest about the limit: an unknown hook shape has no known neutral
            # result, so it is installed as-is and named. Silently pretending it
            # was guarded would be worse than saying it was not.
            self.unguarded.append(name)
            return target

        index = CALLABLE_AT.get(name, 0)

        def install(*args, **kwargs):
            args = list(args)
            if len(args) > index and callable(args[index]):
                if name == "add_wrapper_with_key":
                    args[index] = guard_wrapper(args[index], self._key, self._dropped)
                elif name == "set_model_unet_function_wrapper":
                    args[index] = guard_unet(args[index], self._key, self._dropped,
                                             self._patcher.model_options.get("model_function_wrapper"))
                else:
                    args[index] = guard(args[index], self._key, neutral, self._dropped)
            return target(*args, **kwargs)

        return install

    def guarded(self, fn, neutral):
        """`fn` guarded like any hook installed through this patcher, for hooks
        that go in through a slot rather than a `set_model_*` method (block
        replacements, the attention override). `neutral` has the hook's own
        signature and returns what the call would have without it."""
        if not self._guarding:
            return fn
        return guard(fn, self._key, neutral, self._dropped)

    def __setattr__(self, name, value):
        setattr(self._patcher, name, value)


def clone(model):
    """`model.clone()` that keeps the fault record a modifier node hung on it: ModelPatcher.clone()
    drops unknown attributes, and the sampler would then clear a record nobody else holds."""
    patched = model.clone()
    record = getattr(model, "funpack_dropped", None)
    if record is not None:
        patched.funpack_dropped = record
    return patched