import os
import sys
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
import tqdm
from .imagefunc import *
from comfy.utils import ProgressBar
sys.path.append(os.path.join(os.path.dirname(__file__), 'BiRefNet_v2'))


def get_models():
    model_path = os.path.join(folder_paths.models_dir, 'BiRefNet', 'pth')
    model_ext = [".pth"]
    model_dict = get_files(model_path, model_ext)
    return model_dict

class LS_LoadBiRefNetModel:

    def __init__(self):
        self.birefnet = None
        self.state_dict = None


    @classmethod
    def INPUT_TYPES(s):
        tmp_list = list(get_models().keys())
        model_list = []
        if 'BiRefNet-general-epoch_244.pth' in tmp_list:
            model_list.append('BiRefNet-general-epoch_244.pth')
            tmp_list.remove('BiRefNet-general-epoch_244.pth')
        model_list.extend(tmp_list)

        return {
            "required": {
                "model": (model_list,),
            },
        }

    RETURN_TYPES = ("BIREFNET_MODEL",)
    RETURN_NAMES = ("birefnet_model",)
    FUNCTION = "load_birefnet_model"
    CATEGORY = '😺dzNodes/LayerMask'

    def load_birefnet_model(self, model):
        from .BiRefNet_v2.models.birefnet import BiRefNet
        from .BiRefNet_v2.utils import check_state_dict
        model_dict = get_models()
        self.birefnet = BiRefNet(bb_pretrained=False)
        self.state_dict = torch.load(model_dict[model], map_location='cpu', weights_only=True)
        self.state_dict = check_state_dict(self.state_dict)
        self.birefnet.load_state_dict(self.state_dict)
        return (self.birefnet,)

class LS_LoadBiRefNetModelV2:
    def __init__(self):
        self.model = None

    @classmethod
    def INPUT_TYPES(s):
        model_list = list(s.birefnet_model_repos.keys())
        return {
            "required": {
                "version": (model_list,{"default": model_list[0]}),
            },
        }
    
    RETURN_TYPES = ("BIREFNET_MODEL",)
    RETURN_NAMES = ("birefnet_model",)
    FUNCTION = "load_birefnet_model"
    CATEGORY = '😺dzNodes/LayerMask'

    birefnet_model_repos = {
        "BiRefNet-General": "ZhengPeng7/BiRefNet",
        "RMBG-2.0": "briaai/RMBG-2.0"
    }

    def load_birefnet_model(self, version):
        birefnet_path = os.path.join(folder_paths.models_dir, 'BiRefNet')
        os.makedirs(birefnet_path, exist_ok=True)

        model_path = os.path.join(birefnet_path, version)

        if version == "BiRefNet-General":
            old_birefnet_path = os.path.join(birefnet_path, 'pth')
            old_model = "BiRefNet-general-epoch_244.pth"
            old_model_path = os.path.join(old_birefnet_path, old_model)
            if os.path.exists(old_model_path):
                from .BiRefNet_v2.models.birefnet import BiRefNet
                from .BiRefNet_v2.utils import check_state_dict
                self.birefnet = BiRefNet(bb_pretrained=False)
                self.state_dict = torch.load(old_model_path, map_location='cpu', weights_only=True)
                self.state_dict = check_state_dict(self.state_dict)
                self.birefnet.load_state_dict(self.state_dict)
                return (self.birefnet,)
        elif not os.path.exists(model_path):
            log(f"Downloading {version} model...")
            repo_id = self.birefnet_model_repos[version]
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt"])

        self.model = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True)
        return (self.model,)

class LS_BiRefNetUltraV2:

    def __init__(self):
        self.NODE_NAME = 'BiRefNetUltraV2'

    @classmethod
    def INPUT_TYPES(cls):

        method_list = ['VITMatte', 'VITMatte(local)', 'PyMatting', 'GuidedFilter', ]
        device_list = ['cuda', 'cpu']
        return {
            "required": {
                "image": ("IMAGE",),
                "birefnet_model": ("BIREFNET_MODEL",),
                "detail_method": (method_list,),
                "detail_erode": ("INT", {"default": 4, "min": 1, "max": 255, "step": 1}),
                "detail_dilate": ("INT", {"default": 2, "min": 1, "max": 255, "step": 1}),
                "black_point": ("FLOAT", {"default": 0.01, "min": 0.01, "max": 0.98, "step": 0.01, "display": "slider"}),
                "white_point": ("FLOAT", {"default": 0.99, "min": 0.02, "max": 0.99, "step": 0.01, "display": "slider"}),
                "process_detail": ("BOOLEAN", {"default": False}),
                "device": (device_list,),
                "max_megapixels": ("FLOAT", {"default": 2.0, "min": 1, "max": 999, "step": 0.1}),
            },
            "optional": {
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", )
    RETURN_NAMES = ("image", "mask", )
    FUNCTION = "birefnet_ultra_v2"
    CATEGORY = '😺dzNodes/LayerMask'

    def birefnet_ultra_v2(self, image, birefnet_model, detail_method, detail_erode, detail_dilate,
                       black_point, white_point, process_detail, device, max_megapixels):
        ret_images = []
        ret_masks = []
        inference_image_size = (1024, 1024)
        if detail_method == 'VITMatte(local)':
            local_files_only = True
        else:
            local_files_only = False

        torch.set_float32_matmul_precision(['high', 'highest'][0])
        birefnet_model.to(device)
        birefnet_model.eval()

        comfy_pbar = ProgressBar(len(image))
        tqdm_pbar = tqdm(total=len(image), desc="Processing BiRefNet")
        for i in image:
            i = torch.unsqueeze(i, 0)
            orig_image = tensor2pil(i).convert('RGB')

            transform_image = transforms.Compose([
                transforms.Resize(inference_image_size),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])

            inference_image = transform_image(orig_image).unsqueeze(0).to(device)

            # Prediction
            with torch.no_grad():
                preds = birefnet_model(inference_image)[-1].sigmoid().cpu()
            pred = preds[0].squeeze()
            pred_pil = transforms.ToPILImage()(pred)
            _mask = pred_pil.resize(inference_image_size)

            resize_sampler = Image.BILINEAR
            _mask = _mask.resize(orig_image.size, resize_sampler)
            brightness_image = ImageEnhance.Brightness(_mask)
            _mask = brightness_image.enhance(factor=1.08)
            _mask = image2mask(_mask)

            detail_range = detail_erode + detail_dilate

            if process_detail:
                if detail_method == 'GuidedFilter':
                    _mask = guided_filter_alpha(i, _mask, detail_range // 6 + 1)
                    _mask = tensor2pil(histogram_remap(_mask, black_point, white_point))
                elif detail_method == 'PyMatting':
                    _mask = tensor2pil(mask_edge_detail(i, _mask, detail_range // 8 + 1, black_point, white_point))
                else:
                    _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate)
                    _mask = generate_VITMatte(orig_image, _trimap, local_files_only=local_files_only, device=device, max_megapixels=max_megapixels)
                    _mask = tensor2pil(histogram_remap(pil2tensor(_mask), black_point, white_point))
            else:
                _mask = tensor2pil(_mask)

            ret_image = RGB2RGBA(orig_image, _mask.convert('L'))
            ret_images.append(pil2tensor(ret_image))
            ret_masks.append(image2mask(_mask))

            comfy_pbar.update(1)
            tqdm_pbar.update(1)

        log(f"{self.NODE_NAME} Processed {len(ret_masks)} image(s).", message_type='finish')
        return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)

NODE_CLASS_MAPPINGS = {
    "LayerMask: BiRefNetUltraV2": LS_BiRefNetUltraV2,
    "LayerMask: LoadBiRefNetModel": LS_LoadBiRefNetModel,
    "LayerMask: LoadBiRefNetModelV2": LS_LoadBiRefNetModelV2
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LayerMask: BiRefNetUltraV2": "LayerMask: BiRefNet Ultra V2",
    "LayerMask: LoadBiRefNetModel": "LayerMask: Load BiRefNet Model",
    "LayerMask: LoadBiRefNetModelV2": "LayerMask: Load BiRefNet Model V2"
}
