"""ComfyUI Memory Manager — prevent OOM crashes with smart model unloading."""
import gc
import time
import logging
from collections import OrderedDict
from typing import Optional

logger = logging.getLogger("comfyui_memory_manager")

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

try:
    import comfy.model_management
    COMFY_AVAILABLE = True
except ImportError:
    COMFY_AVAILABLE = False


class LRUModelCache:
    """Least Recently Used cache for GPU models."""

    def __init__(self, max_size_gb: float = 6.0):
        self.max_size_bytes = int(max_size_gb * 1024**3)
        self.cache: OrderedDict = OrderedDict()
        self.current_size = 0

    def get(self, key: str):
        if key in self.cache:
            self.cache.move_to_end(key)
            return self.cache[key]
        return None

    def put(self, key: str, model, size_bytes: int):
        if key in self.cache:
            self.cache.move_to_end(key)
            return
        self.cache[key] = {"model": model, "size": size_bytes}
        self.current_size += size_bytes
        self._evict()

    def _evict(self):
        while self.current_size > self.max_size_bytes and self.cache:
            key, entry = self.cache.popitem(last=False)
            self.current_size -= entry["size"]
            logger.info(f"Evicted model: {key} ({entry['size'] / 1024**2:.0f} MB)")
            del entry["model"]
            gc.collect()
            if TORCH_AVAILABLE and torch.cuda.is_available():
                torch.cuda.empty_cache()

    def clear(self):
        self.cache.clear()
        self.current_size = 0
        gc.collect()
        if TORCH_AVAILABLE and torch.cuda.is_available():
            torch.cuda.empty_cache()

    @property
    def usage_bytes(self):
        return self.current_size

    @property
    def usage_gb(self):
        return self.current_size / 1024**3


def get_vram_info() -> dict:
    """Get current GPU memory information."""
    if not TORCH_AVAILABLE or not torch.cuda.is_available():
        return {"available": False, "total_mb": 0, "used_mb": 0, "free_mb": 0}

    total = torch.cuda.get_device_properties(0).total_mem / 1024**2
    reserved = torch.cuda.memory_reserved(0) / 1024**2
    allocated = torch.cuda.memory_allocated(0) / 1024**2
    free = total - reserved

    return {
        "available": True,
        "total_mb": round(total),
        "reserved_mb": round(reserved),
        "allocated_mb": round(allocated),
        "free_mb": round(free),
        "utilization": round(allocated / total * 100, 1),
    }


def force_cleanup():
    """Force garbage collection and CUDA cache clear."""
    gc.collect()
    if TORCH_AVAILABLE and torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


def estimate_model_size(model) -> Optional[int]:
    """Estimate model VRAM usage in bytes."""
    if not TORCH_AVAILABLE:
        return None
    try:
        if hasattr(model, "model"):
            return sum(p.numel() * p.element_size() for p in model.model.parameters())
        return sum(p.numel() * p.element_size() for p in model.parameters())
    except Exception:
        return None


# Global cache instance
_model_cache = LRUModelCache(max_size_gb=6.0)


def get_cache() -> LRUModelCache:
    return _model_cache


def set_cache_size(gb: float):
    _model_cache.max_size_bytes = int(gb * 1024**3)