import psutil
import ctypes
from ctypes import wintypes
import time
import platform
import gc
from server import PromptServer
import comfy.model_management

try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False

class AnyType(str):
    def __eq__(self, _) -> bool:
        return True
    def __ne__(self, __value: object) -> bool:
        return False

any = AnyType("*")


def safe_int(value, default=0):
    if isinstance(value, int):
        return value
    try:
        return int(float(str(value).strip()))
    except (TypeError, ValueError):
        return default


def safe_bool(value, default=False):
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return bool(value)
    if isinstance(value, str):
        return value.strip().lower() in ("true", "1", "yes")
    return default


def check_memory_critical(threshold_mb=500):
    try:
        memory = psutil.virtual_memory()
        available_mb = memory.available / (1024 * 1024)
        return available_mb < threshold_mb, available_mb
    except:
        return False, 0


def force_memory_cleanup():
    gc.collect()
    gc.collect()
    gc.collect()
    if HAS_TORCH and torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    try:
        comfy.model_management.soft_empty_cache()
    except:
        pass


class VRAMCleanup:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "batch_index": ("INT", {
                    "default": 60, "min": 0, "step": 1,
                    "label": "批次索引"
                }),
                "batch_clear": ("INT", {
                    "default": 60, "min": 1, "step": 1,
                    "label": "清理间隔"
                }),
                "bypass_all": ("BOOLEAN", {"default": False, "label": "忽略全部"}),
                "offload_model": ("BOOLEAN", {"default": True}),
                "offload_cache": ("BOOLEAN", {"default": True}),
                "critical_threshold_mb": ("INT", {
                    "default": 500, "min": 100, "max": 8192, "step": 100,
                    "label": "危险阈值(MB)"
                }),
            },
            "optional": {
                "anything": (any, {}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "extra_pnginfo": "EXTRA_PNGINFO",
            }
        }

    RETURN_TYPES = (any, "INT",)
    RETURN_NAMES = ("anything", "batch_index",)
    OUTPUT_NODE = True
    FUNCTION = "empty_cache"
    CATEGORY = "Memory Management"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float(time.time())

    def empty_cache(self, bypass_all, offload_model, offload_cache, batch_index, batch_clear, critical_threshold_mb=500, anything=None, unique_id=None, extra_pnginfo=None):
        try:
            bypass_all = safe_bool(bypass_all, False)
            offload_model = safe_bool(offload_model, True)
            offload_cache = safe_bool(offload_cache, True)
            batch_index = safe_int(batch_index, 0)
            batch_clear = safe_int(batch_clear, 1)
            critical_threshold_mb = safe_int(critical_threshold_mb, 500)

            if batch_clear <= 0:
                batch_clear = 1

            is_critical, available_mb = check_memory_critical(critical_threshold_mb)
            should_execute = (batch_index % batch_clear == 0)

            if bypass_all and not is_critical:
                print(f"VRAM清理跳过 (bypass_all=True, 可用: {available_mb:.0f}MB)")
            elif not should_execute and not is_critical:
                print(f"VRAM清理跳过 (index={batch_index} % batch={batch_clear} = {batch_index % batch_clear}, 可用: {available_mb:.0f}MB)")
            else:
                if is_critical:
                    print(f"⚠️ 内存危险! 可用仅 {available_mb:.0f}MB < {critical_threshold_mb}MB, 强制清理!")

                if offload_model:
                    comfy.model_management.unload_all_models()

                if offload_cache:
                    force_memory_cleanup()
                    PromptServer.instance.prompt_queue.set_flag("free_memory", True)

                _, after_mb = check_memory_critical(critical_threshold_mb)
                print(f"VRAM清理完成 [卸载模型: {offload_model}, 清空缓存: {offload_cache}, index: {batch_index}, 可用: {available_mb:.0f}MB -> {after_mb:.0f}MB]")

        except Exception as e:
            print(f"VRAM清理失败: {type(e).__name__}: {e}")
            try:
                force_memory_cleanup()
            except:
                pass

        return (anything, batch_index,)


class RAMCleanup:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "batch_index": ("INT", {
                    "default": 30, "min": 0, "step": 1,
                    "label": "当前索引"
                }),
                "batch_clear": ("INT", {
                    "default": 30, "min": 1, "step": 1,
                    "label": "批次间隔"
                }),
                "bypass_all": ("BOOLEAN", {"default": False, "label": "忽略全部"}),
                "clean_file_cache": ("BOOLEAN", {"default": True, "label": "清理文件缓存"}),
                "clean_processes": ("BOOLEAN", {"default": True, "label": "清理进程内存"}),
                "clean_dlls": ("BOOLEAN", {"default": True, "label": "清理未使用DLL"}),
                "retry_times": ("INT", {
                    "default": 3, "min": 1, "max": 10, "step": 1,
                    "label": "重试次数"
                }),
                "critical_threshold_mb": ("INT", {
                    "default": 500, "min": 100, "max": 8192, "step": 100,
                    "label": "危险阈值(MB)"
                }),
            },
            "optional": {
                "anything": (any, {}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "extra_pnginfo": "EXTRA_PNGINFO",
            }
        }

    # ★★★ 两个输出：anything + batch_index ★★★
    RETURN_TYPES = (any, "INT",)
    RETURN_NAMES = ("anything", "batch_index",)
    OUTPUT_NODE = True
    FUNCTION = "clean_ram"
    CATEGORY = "Memory Management"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float(time.time())

    def get_ram_usage(self):
        memory = psutil.virtual_memory()
        return memory.percent, memory.available / (1024 * 1024)

    def clean_ram(self, bypass_all, clean_file_cache, clean_processes, clean_dlls, retry_times, batch_index, batch_clear, critical_threshold_mb=500, anything=None, unique_id=None, extra_pnginfo=None):
        try:
            bypass_all = safe_bool(bypass_all, False)
            clean_file_cache = safe_bool(clean_file_cache, True)
            clean_processes = safe_bool(clean_processes, True)
            clean_dlls = safe_bool(clean_dlls, True)
            retry_times = safe_int(retry_times, 3)
            batch_index = safe_int(batch_index, 0)
            batch_clear = safe_int(batch_clear, 1)
            critical_threshold_mb = safe_int(critical_threshold_mb, 500)

            if batch_clear <= 0:
                batch_clear = 1
            if retry_times <= 0:
                retry_times = 1

            before_usage, before_available = self.get_ram_usage()
            system = platform.system()

            is_critical, available_mb = check_memory_critical(critical_threshold_mb)
            should_execute = (batch_index % batch_clear == 0)

            if bypass_all and not is_critical:
                print(f"RAM清理跳过 (bypass_all=True, 可用: {available_mb:.0f}MB)")
            elif not should_execute and not is_critical:
                print(f"RAM清理跳过 (index={batch_index} % batch={batch_clear} = {batch_index % batch_clear}, 可用: {available_mb:.0f}MB)")
            else:
                if is_critical:
                    print(f"⚠️ 内存危险! 可用仅 {available_mb:.0f}MB < {critical_threshold_mb}MB, 强制清理!")

                force_memory_cleanup()

                for attempt in range(retry_times):
                    if clean_file_cache:
                        try:
                            if system == "Windows":
                                ctypes.windll.kernel32.SetSystemFileCacheSize(-1, -1, 0)
                            elif system == "Linux":
                                libc = ctypes.CDLL("libc.so.6")
                                libc.malloc_trim(0)
                        except:
                            pass

                    if clean_processes:
                        if system == "Windows":
                            for process in psutil.process_iter(['pid', 'name']):
                                try:
                                    handle = ctypes.windll.kernel32.OpenProcess(
                                        wintypes.DWORD(0x001F0FFF),
                                        wintypes.BOOL(False),
                                        wintypes.DWORD(process.info['pid'])
                                    )
                                    ctypes.windll.psapi.EmptyWorkingSet(handle)
                                    ctypes.windll.kernel32.CloseHandle(handle)
                                except:
                                    continue

                    if clean_dlls:
                        try:
                            if system == "Windows":
                                ctypes.windll.kernel32.SetProcessWorkingSetSize(-1, -1, -1)
                        except:
                            pass

                    time.sleep(1)

                    _, current_mb = check_memory_critical(critical_threshold_mb)
                    if current_mb > critical_threshold_mb * 2:
                        print(f"  第{attempt + 1}次清理后内存充足 ({current_mb:.0f}MB), 停止重试")
                        break

                after_usage, after_available = self.get_ram_usage()
                freed_mb = after_available - before_available
                print(f"RAM清理完成 [{before_usage:.1f}% -> {after_usage:.1f}%, 释放: {freed_mb:.0f}MB, index: {batch_index}, batch: {batch_clear}]")

        except Exception as e:
            print(f"RAM清理失败: {type(e).__name__}: {e}")
            print(f"  调试 -> batch_index={repr(batch_index)}(type={type(batch_index).__name__}), batch_clear={repr(batch_clear)}(type={type(batch_clear).__name__})")
            try:
                force_memory_cleanup()
            except:
                pass

        # ★★★ 返回 anything + batch_index ★★★
        return (anything, batch_index,)


NODE_CLASS_MAPPINGS = {
    "DocKr.VRAM-Cleanup": VRAMCleanup,
    "DocKr.RAM-Cleanup": RAMCleanup,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DocKr.VRAM-Cleanup": "DocKr / VRAM-Cleanup",
    "DocKr.RAM-Cleanup": "DocKr / RAM-Cleanup",
}