# Shared utility helpers
import os
import gc
from typing import Union, Optional
import numpy as np
import torch
from PIL import Image, ImageFilter
from torch.hub import download_url_to_file
from comfy.utils import common_upscale
import comfy.model_management as mm
import folder_paths


def get_device():
    """Return the torch device ComfyUI selected (CUDA, MPS, or CPU)."""
    return mm.get_torch_device()


def clean_vram():
    """Safely clear cached VRAM across CUDA, MPS, and CPU backends."""
    gc.collect()
    mm.soft_empty_cache()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        if hasattr(torch.cuda, "ipc_collect"):
            torch.cuda.ipc_collect()
    elif hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
        try:
            torch.mps.empty_cache()
        except Exception:
            pass


def unload_model(loader_or_model, attr_name="model"):
    """Unload a model from memory and release VRAM.

    Accepts either:
    1. A loader object that owns a model attribute (e.g. loader, attr_name='model')
    2. A raw torch.nn.Module instance directly.
    """
    if hasattr(loader_or_model, attr_name):
        model = getattr(loader_or_model, attr_name)
        if model is not None:
            if hasattr(model, "to"):
                try:
                    model.to("cpu")
                except Exception:
                    pass
            setattr(loader_or_model, attr_name, None)
            if hasattr(loader_or_model, "current_model_version"):
                loader_or_model.current_model_version = None
    elif isinstance(loader_or_model, torch.nn.Module):
        try:
            loader_or_model.to("cpu")
        except Exception:
            pass
        del loader_or_model

    clean_vram()


def patch_transformers_bert():
    """Ensure BertModel and BertModelWarper compatibility for GroundingDINO on transformers >= 5.0."""
    try:
        def _get_head_mask(self, head_mask=None, num_hidden_layers=12, is_attention_chunked=False):
            if head_mask is not None:
                if hasattr(head_mask, "dim"):
                    if head_mask.dim() == 1:
                        head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                        head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
                    elif head_mask.dim() == 2:
                        head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
                if is_attention_chunked and hasattr(head_mask, "unsqueeze"):
                    head_mask = head_mask.unsqueeze(-1)
                return head_mask
            return [None] * num_hidden_layers

        def _adapted_get_extended_attention_mask(self, attention_mask, input_shape=None, *args, **kwargs):
            device = kwargs.pop("device", None)
            dtype = kwargs.pop("dtype", None)
            if len(args) >= 1:
                arg = args[0]
                if isinstance(arg, torch.device) or (isinstance(arg, str) and str(arg).startswith(("cuda", "cpu", "mps"))):
                    device = arg
                elif isinstance(arg, torch.dtype):
                    dtype = arg
            if len(args) >= 2:
                arg2 = args[1]
                if isinstance(arg2, torch.dtype):
                    dtype = arg2
                elif isinstance(arg2, (torch.device, str)):
                    device = arg2

            if dtype is None:
                dtype = getattr(self, "dtype", torch.float32)

            if attention_mask.dim() == 3:
                extended_attention_mask = attention_mask[:, None, :, :]
            elif attention_mask.dim() == 2:
                extended_attention_mask = attention_mask[:, None, None, :]
            else:
                raise ValueError(f"Wrong shape for attention_mask (shape {attention_mask.shape})")

            if device is not None:
                extended_attention_mask = extended_attention_mask.to(device=device)
            if dtype is not None:
                extended_attention_mask = extended_attention_mask.to(dtype=dtype)

            extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
            return extended_attention_mask

        # 1. Patch BertModelWarper directly in GroundingDINO
        try:
            import groundingdino.models.GroundingDINO.bertwarper as bw
            if not getattr(bw, "_rmbg_patched", False):
                orig_init = bw.BertModelWarper.__init__
                def new_init(self, bert_model):
                    if not hasattr(bert_model, "get_head_mask"):
                        bert_model.get_head_mask = lambda *a, **k: _get_head_mask(bert_model, *a, **k)
                    bert_model.get_extended_attention_mask = lambda *a, **k: _adapted_get_extended_attention_mask(bert_model, *a, **k)
                    orig_init(self, bert_model)
                    if not hasattr(self, "get_head_mask") or self.get_head_mask is None:
                        self.get_head_mask = lambda *a, **k: _get_head_mask(self, *a, **k)
                    self.get_extended_attention_mask = lambda *a, **k: _adapted_get_extended_attention_mask(self, *a, **k)
                bw.BertModelWarper.__init__ = new_init
                bw._rmbg_patched = True
        except Exception:
            pass

        # 2. Patch BertModel in transformers
        try:
            import transformers
            if hasattr(transformers, "BertModel"):
                if not hasattr(transformers.BertModel, "get_head_mask"):
                    transformers.BertModel.get_head_mask = _get_head_mask
                transformers.BertModel.get_extended_attention_mask = _adapted_get_extended_attention_mask
            if hasattr(transformers, "PreTrainedModel"):
                if not hasattr(transformers.PreTrainedModel, "get_head_mask"):
                    transformers.PreTrainedModel.get_head_mask = _get_head_mask
                transformers.PreTrainedModel.get_extended_attention_mask = _adapted_get_extended_attention_mask
            from transformers.models.bert.modeling_bert import BertModel
            if not hasattr(BertModel, "get_head_mask"):
                BertModel.get_head_mask = _get_head_mask
            BertModel.get_extended_attention_mask = _adapted_get_extended_attention_mask
        except Exception:
            pass

    except Exception:
        pass


def tensor2pil(image: torch.Tensor) -> Image.Image:
    return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))


def pil2tensor(image: Image.Image) -> torch.Tensor:
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)


def pil2mask(image: Image.Image) -> torch.Tensor:
    return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0).unsqueeze(0)


def mask2image(mask: torch.Tensor) -> Image.Image:
    if len(mask.shape) == 2:
        mask = mask.unsqueeze(0)
    return tensor2pil(mask)


def image2mask(image: Image.Image) -> torch.Tensor:
    return torch.from_numpy(np.array(image.convert("L")).astype(np.float32) / 255.0)


def RGB2RGBA(image: Image.Image, mask: Union[Image.Image, torch.Tensor]) -> Image.Image:
    if isinstance(mask, torch.Tensor):
        mask = mask2image(mask)
    if mask.size != image.size:
        mask = mask.resize(image.size, Image.Resampling.LANCZOS)
    return Image.merge('RGBA', (*image.convert('RGB').split(), mask.convert('L')))


def process_mask(mask_image: Image.Image, invert_output: bool = False, mask_blur: int = 0, mask_offset: int = 0) -> Image.Image:
    if invert_output:
        mask_np = np.array(mask_image)
        mask_image = Image.fromarray(255 - mask_np)
    if mask_blur > 0:
        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=mask_blur))
    if mask_offset != 0:
        filter_type = ImageFilter.MaxFilter if mask_offset > 0 else ImageFilter.MinFilter
        size = abs(mask_offset) * 2 + 1
        for _ in range(abs(mask_offset)):
            mask_image = mask_image.filter(filter_type(size))
    return mask_image


def apply_background_color(image: Image.Image, mask_image: Image.Image, background: str = "Alpha", background_color: str = "#222222") -> Image.Image:
    rgba_image = image.copy().convert('RGBA')
    rgba_image.putalpha(mask_image.convert('L'))
    if background == "Color":
        def hex_to_rgba(hex_color):
            hex_color = hex_color.lstrip('#')
            r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
            return (r, g, b, 255)
        rgba = hex_to_rgba(background_color)
        bg_image = Image.new('RGBA', image.size, rgba)
        composite_image = Image.alpha_composite(bg_image, rgba_image)
        return composite_image.convert('RGB')
    return rgba_image


def get_or_download_model_file(filename: str, url: str, dirname: str = "rmbg") -> str:
    local_path = folder_paths.get_full_path(dirname, filename)
    if local_path:
        return local_path
    folder = os.path.join(folder_paths.models_dir, dirname)
    os.makedirs(folder, exist_ok=True)
    local_path = os.path.join(folder, filename)
    if not os.path.exists(local_path):
        print(f"Downloading {filename} from {url} ...")
        download_url_to_file(url, local_path)
    return local_path


def resize_image(img: Image.Image, width: int, height: int) -> Image.Image:
    return img.resize((width, height), resample=Image.LANCZOS)


def blend_overlay(img_1: Image.Image, img_2: Image.Image) -> Image.Image:
    arr1 = np.array(img_1).astype(float) / 255.0
    arr2 = np.array(img_2).astype(float) / 255.0
    mask = arr2 < 0.5
    result = np.zeros_like(arr1)
    result[mask] = 2 * arr1[mask] * arr2[mask]
    result[~mask] = 1 - 2 * (1 - arr1[~mask]) * (1 - arr2[~mask])
    return Image.fromarray(np.clip(result * 255, 0, 255).astype(np.uint8))


def fill_mask(width: int, height: int, mask: Image.Image, box=(0, 0), color=0) -> Image.Image:
    bg = Image.new("L", (width, height), color)
    bg.paste(mask, box, mask)
    return bg


def empty_image(width: int, height: int, batch_size: int = 1) -> torch.Tensor:
    return torch.zeros([batch_size, height, width, 3])


def upscale_mask(mask: torch.Tensor, width: int, height: int) -> torch.Tensor:
    if mask.ndim == 3:
        mask = mask.unsqueeze(1)
    mask = common_upscale(mask, width, height, "bicubic", "disabled")
    mask = mask.squeeze(1)
    return mask


def extract_alpha_mask(image: torch.Tensor) -> torch.Tensor:
    alpha = image[..., 3]
    if alpha.max() > 1.0:
        alpha = alpha / 255.0
    if len(alpha.shape) == 4:
        alpha = alpha[:, :, :, 0]
    return alpha.unsqueeze(1) if alpha.ndim == 3 else alpha


def ensure_mask_shape(mask: torch.Tensor | None) -> torch.Tensor | None:
    if mask is None:
        return None
    if mask.ndim == 2:
        return mask.unsqueeze(0)
    if mask.ndim == 4 and mask.shape[1] == 1:
        return mask.squeeze(1)
    return mask


COLOR_PRESETS = {
    "black": "#000000",
    "white": "#FFFFFF",
    "red": "#FF0000",
    "green": "#00FF00",
    "blue": "#0000FF",
    "yellow": "#FFFF00",
    "cyan": "#00FFFF",
    "magenta": "#FF00FF",
    "gray": "#808080",
    "silver": "#C0C0C0",
    "maroon": "#800000",
    "olive": "#808000",
    "purple": "#800080",
    "teal": "#008080",
    "navy": "#000080",
    "orange": "#FFA500",
    "pink": "#FFC0CB",
    "brown": "#A52A2A",
    "violet": "#EE82EE",
    "indigo": "#4B0082",
    "light_gray": "#D3D3D3",
    "dark_gray": "#A9A9A9",
    "light_blue": "#ADD8E6",
    "dark_blue": "#00008B",
    "light_green": "#90EE90",
    "dark_green": "#006400",
}

ASPECT_RATIOS = {
    # Square
    "1:1 (Square)":               (1, 1),

    # Landscape
    "5:4 (Landscape)":            (5, 4),
    "4:3 (Standard)":             (4, 3),
    "3:2 (Photo)":                (3, 2),
    "16:10 (Landscape)":          (16, 10),
    "16:9 (Widescreen)":          (16, 9),
    "1.85:1 (Film)":              (185, 100),
    "2:1 (Univisium)":            (2, 1),
    "21:9 (Ultrawide)":           (21, 9),
    "2.35:1 (Anamorphic)":        (235, 100),
    "2.39:1 (Panavision)":        (239, 100),
    "32:9 (Super Ultrawide)":     (32, 9),
    "3:1 (Panorama)":             (3, 1),

    # Portrait (mirrors of the landscape ratios)
    "4:5 (Portrait)":             (4, 5),
    "3:4 (Portrait Standard)":    (3, 4),
    "2:3 (Portrait Photo)":       (2, 3),
    "1:2 (Portrait Univisium)":   (1, 2),
    "10:16 (Portrait)":           (10, 16),
    "9:16 (Portrait Widescreen)": (9, 16),
    "9:21 (Portrait Ultrawide)":  (9, 21),
    "1:3 (Portrait Panorama)":    (1, 3),
}

def color_format(color: str) -> str:
    if not color:
        return ""

    color = color.strip().upper()
    if not color.startswith("#"):
        color = f"#{color}"

    color = color[1:]
    if len(color) == 3:
        r, g, b = color[0], color[1], color[2]
        return f"#{r}{r}{g}{g}{b}{b}"
    if len(color) < 6:
        raise ValueError(f"Invalid color format: {color}")
    if len(color) > 6:
        color = color[:6]

    return f"#{color}"


def is_transformers_supported(min_version=(5, 0, 0)) -> bool:
    try:
        import transformers
        parts = []
        for p in transformers.__version__.split(".")[:3]:
            digits = "".join(c for c in p if c.isdigit())
            parts.append(int(digits) if digits else 0)
        while len(parts) < 3:
            parts.append(0)
        return tuple(parts[:3]) >= min_version
    except Exception:
        return False