import argparse
import hashlib
import os
import sys
import re

import folder_paths

def add_search_paths():
    # Set the configs folder
    custom_nodes_folders, _ = folder_paths.folder_names_and_paths["custom_nodes"]

    cartoon_seg_folders = {}

    # custom_nodes folder is an array.  This code assumes that there be only one instance
    custom_node_folder = custom_nodes_folders[0].rstrip("/").rstrip("\\")
    # Do Not Add to Lookups.  Force custom keys as way of tracking unique loads/calls
    cartoon_segmentation_root = f"{custom_node_folder}/ComfyUI_CartoonSegmentation/CartoonSegmentation"
    cartoon_seg_config_path = f"{cartoon_segmentation_root}/configs"
    model_path_ani_inst_ckpt = f"{cartoon_segmentation_root}/models/AnimeInstanceSegmentation"
    # model_path_ani_seg_ckpt = f"{cartoon_segmentation_root}/models/anime-seg"
    leres_path = f"{cartoon_segmentation_root}/models/leres"
    depth_modules_path = f"{cartoon_segmentation_root}/depth_modules"

    if "cartoon_segmentation" not in folder_paths.folder_names_and_paths:
        folder_paths.folder_names_and_paths["cartoon_segmentation"] = cartoon_seg_folders
    else:
        cartoon_seg_folders = folder_paths.folder_names_and_paths["cartoon_segs_configs"]

    # End Goal is to have three new entries for Paths.  Configs and models in ComfyUI searchable structures and
    # a third key that contains a dictionary for folders that have hard references

    def check_report_path(path, variable_target):
        if path_exists := not os.path.exists(path):
            print(f"Warning Excepted Path {variable_target} "
                  f"does not exist when loading ComfyUI_CartoonSegmentation\n"
                  f"Path::{path}")
        return path_exists

    if not check_report_path(custom_node_folder, "custom node root folder"):
        cartoon_seg_folders["custom_node_root"] = custom_node_folder

    if not check_report_path(cartoon_seg_config_path, "Cartoon Segmentation Configs"):
        folder_paths.folder_names_and_paths["cartoon_seg_configs"] = ([custom_node_folder], {'.yaml'})

    if not check_report_path(model_path_ani_inst_ckpt, "Cartoon Segmentation checkpoints AnimeInstanceSegmentation"):
        folder_paths.folder_names_and_paths["cartoon_seg_ckpt"] = ([model_path_ani_inst_ckpt], {'.ckpt', '.safetensors'})
        cartoon_seg_folders["anime_instance_path"] = model_path_ani_inst_ckpt
        # Add Models to the sys.path.  A call to Config.loadString in AnimeInsSeg external lib mmengine needs help
        sys.path.append(model_path_ani_inst_ckpt)

    # if not check_report_path(model_path_ani_seg_ckpt, "Cartoon Segmentation anime-seg checkpoints"):
    #     cartoon_seg_folders["anime_seg_checkpoints"] = model_path_ani_seg_ckpt


    def directory_path_to_module_path(directory_path):
        root_split = re.split(r"/|\\", directory_path)
        # There is a calls that do not have a full path, but a relative path for import
        comfy_root_length = len(re.split(r"/|\\", folder_paths.base_path))
        base_dot_path = ".".join(root_split[comfy_root_length:])
        return base_dot_path

    if not check_report_path(leres_path, "leres_path"):
        cartoon_seg_folders["leres_import_path"] = directory_path_to_module_path(leres_path)
        cartoon_seg_folders["leres_model_path"] = leres_path

    if not check_report_path(depth_modules_path, "depth_modules"):
        cartoon_seg_folders["depth_modules"] = directory_path_to_module_path(depth_modules_path)


add_search_paths()

class KenBurnsConfigLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "config_file": (folder_paths.get_filename_list("cartoon_seg_configs"), ),
            "model": (folder_paths.get_filename_list("cartoon_seg_ckpt"), {"default": "rtmdetl_e60.ckpt"} ),
            "image_path": ("image_path", {"display": "Input Image File Name"}),
            "verbose": (["Yes", "No"], {"display": "Verbose Logging", "default": "No"}),
            },
        }

    RETURN_TYPES = ('kb_config', )
    FUNCTION = 'load_kb_config'
    CATEGORY = "CartoonSegmentation"

    def load_kb_config(self, config_file, model, image_path, verbose):
        cfg_full_path = os.path.join(folder_paths.folder_names_and_paths["cartoon_seg_configs"][0][0], config_file)
        model_full_path = os.path.join(folder_paths.folder_names_and_paths["cartoon_seg_ckpt"][0][0], model)
        return_value = cfg_full_path, model_full_path, image_path, verbose == "Yes"
        return return_value,

class LoadImageFilename:
    @classmethod
    def INPUT_TYPES(s):
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        return {"required":
                    {"image": (sorted(files), {"image_upload": True})},
                }

    CATEGORY = "image"

    RETURN_TYPES = ("image_path", )
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)

        return (image_path, )

    @classmethod
    def IS_CHANGED(s, image):
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()

    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        return True