"""Force memory cleanup nodes for ComfyUI.

This extension deliberately asks ComfyUI's prompt worker to clean up *after*
the current prompt has finished.  Cleaning synchronously inside a node is too
early because the executor still owns references to the current graph.
"""

from __future__ import annotations

import logging
from typing import Any

from server import PromptServer


LOGGER = logging.getLogger("ComfyUI-Force-Cleanup")


class ForceCleanupAfterPrompt:
    """Request a model/cache cleanup after the current prompt completes."""

    CATEGORY = "utils/memory"
    FUNCTION = "request_cleanup"
    RETURN_TYPES = ()
    OUTPUT_NODE = True
    DESCRIPTION = (
        "Requests cleanup after the entire prompt finishes. Full cleanup unloads "
        "models, clears ComfyUI's execution cache, runs Python GC, and releases "
        "the PyTorch device allocator cache."
    )

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return {
            "required": {
                "cleanup_level": (
                    [
                        "full: models + execution cache + allocator",
                        "models: unload models + allocator",
                    ],
                    {
                        "default": "full: models + execution cache + allocator",
                        "tooltip": (
                            "Full frees the most memory but the next run must "
                            "reload/recompute everything."
                        ),
                    },
                ),
            }
        }

    @classmethod
    def IS_CHANGED(cls, **_: Any) -> float:
        # Output nodes can otherwise be satisfied from the execution cache.
        # NaN never compares equal, so the cleanup request is made on every run.
        return float("nan")

    def request_cleanup(self, cleanup_level: str):
        prompt_server = getattr(PromptServer, "instance", None)
        prompt_queue = getattr(prompt_server, "prompt_queue", None)
        if prompt_queue is None:
            raise RuntimeError(
                "ComfyUI prompt queue is unavailable; cleanup was not scheduled."
            )

        if cleanup_level.startswith("full:"):
            # ComfyUI's worker handles this after execution:
            # unload_all_models() -> executor.reset() -> gc.collect()
            # -> soft_empty_cache().
            prompt_queue.set_flag("free_memory", True)
            message = (
                "Full cleanup scheduled after this prompt: models, execution "
                "cache, Python GC, and device allocator cache."
            )
        else:
            # The worker unloads models, runs GC, and empties the device allocator,
            # while retaining the executor's reusable graph cache.
            prompt_queue.set_flag("unload_models", True)
            message = (
                "Model cleanup scheduled after this prompt; reusable execution "
                "cache will be kept."
            )

        LOGGER.info(message)
        return {"ui": {"text": [message]}, "result": ()}


NODE_CLASS_MAPPINGS = {
    "ForceCleanupAfterPrompt": ForceCleanupAfterPrompt,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ForceCleanupAfterPrompt": "Force Cleanup After Prompt",
}

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]