"""
VRAMFree Node for ComfyUI
MIT License - Independent implementation without GPL code reuse
"""

import logging

ANY_SOCKET_TYPES = "IMAGE,LATENT,MODEL,CONDITIONING,CLIP,VAE,MASK,INT,FLOAT,STRING,BOOLEAN"

try:
    from .memory_utils import MemoryManager
except ImportError:
    from memory_utils import MemoryManager


class VRAMFreeNode:
    """ComfyUI custom node that releases memory by unloading models and clearing caches."""

    def __init__(self):
        self.memory_manager = MemoryManager()

    @classmethod
    def INPUT_TYPES(cls):
        """Define node input configuration."""
        return {
            "required": {},
            "optional": {
                "Any": (ANY_SOCKET_TYPES,),
                "MinFreeVRAMGB": ("FLOAT", {
                    "default": 0.0,
                    "min": 0.0,
                    "step": 0.1,
                    "tooltip": "Skip cleanup when free VRAM is already above this threshold (GB)",
                }),
            },
        }

    RETURN_TYPES = (ANY_SOCKET_TYPES,)
    RETURN_NAMES = ("Bypass(Relay)",)

    FUNCTION = "execute_cleanup"
    CATEGORY = "u5"
    OUTPUT_NODE = True
    DESCRIPTION = "Unload models and free memory without disrupting workflow"

    def execute_cleanup(self, Any=None, MinFreeVRAMGB=0.0, **legacy_kwargs):
        """Run the memory cleanup routine and return status details."""
        legacy_detail_log = legacy_kwargs.pop("detail_log", None)
        legacy_verbose = legacy_kwargs.pop("verbose", None)

        # Legacy compatibility for removed flags (always True now)
        legacy_kwargs.pop("unload_models", None)
        legacy_kwargs.pop("clear_cache", None)
        legacy_kwargs.pop("force_gc", None)
        legacy_min_vram = legacy_kwargs.pop("min_free_vram_gb", None)
        if legacy_min_vram is not None:
            MinFreeVRAMGB = legacy_min_vram

        detail_debug_log = bool(legacy_detail_log or legacy_verbose)

        try:
            initial_status = self.memory_manager.get_memory_status()
            initial_summary = self.memory_manager.get_memory_summary()

            if detail_debug_log:
                print("\n=== VRAMFree: Starting Cleanup ===")
                print(f"Initial Memory Status:\n{initial_summary}")
                print("Options: unload_models=True, clear_cache=True, force_gc=True")

            result = self.memory_manager.cleanup_memory(
                unload_models=True,
                clear_cache=True,
                force_gc=True,
                min_free_vram_gb=MinFreeVRAMGB
            )

            final_status = result.get("final_status") or {}
            freed_vram = float(result.get("freed_vram_gb", 0.0))
            freed_ram = float(result.get("freed_ram_gb", 0.0))
            total_freed = max(0.0, freed_vram + freed_ram)

            status_parts = ["Models", "Cache", "GC"]
            actions_label = ", ".join(status_parts)
            freed_str = self.memory_manager.format_memory_size(total_freed)
            status_line = f"Cleared: {actions_label} | Freed: {freed_str}" if status_parts else "No actions performed"

            final_summary = self.memory_manager.get_memory_summary(result.get("final_status"))

            report_lines = [
                "=== Memory Cleanup Report ===",
                "",
                f"Status: {'SUCCESS' if result.get('success') else 'FAILED'}",
                f"Actions: {status_line}",
                "",
                "--- Memory Changes ---",
                f"VRAM Freed: {self.memory_manager.format_memory_size(freed_vram)}",
                f"RAM Freed: {self.memory_manager.format_memory_size(freed_ram)}",
                f"Total Freed: {self.memory_manager.format_memory_size(total_freed)}",
                "",
                "--- Current Memory Status ---",
                final_summary,
            ]

            if detail_debug_log:
                report_lines.append("")
                report_lines.append("--- Detailed Actions ---")
                for action in result.get("actions_performed", []):
                    report_lines.append(f"  - {action}")

            if not result.get("success"):
                report_lines.append("")
                report_lines.append(f"Error: {result.get('message', 'Unknown error')}")

            report = "\n".join(report_lines)

            logging.info(f"VRAMFree: {status_line}")

            if detail_debug_log:
                print(report)

            return (Any,)

        except Exception as exc:  # noqa: BLE001
            status_line = f"Exception: {exc}"
            report_lines = [
                "=== Memory Cleanup Report ===",
                "",
                "Status: EXCEPTION",
                status_line,
                "",
                "Please check the console for detailed error information.",
            ]
            report = "\n".join(report_lines)

            logging.error(f"VRAMFree: {status_line}")

            if detail_debug_log:
                print(f"\nEXCEPTION: {exc}")
                import traceback
                traceback.print_exc()

            return (Any,)

    @classmethod
    def IS_CHANGED(cls, **_kwargs):
        """Force reevaluation so the node executes every time."""
        return float("NaN")