import yaml  # type: ignore
import os
import folder_paths  # type: ignore
import re
import time
import numpy as np  # type: ignore
import torch  # type: ignore
from PIL import Image  # type: ignore
import json
import base64
import io


class FileHandler:
    def __init__(self, filename):
        self.filename = filename
        self.lines = None

    def read_line(self, line_index):
        """Read a specific line from the file by its index (0-based)."""
        with open(self.filename, "r", encoding="utf-8") as file:
            lines = file.readlines()
            lines = [line for line in lines if line.strip()]  # Ignore empty lines
            if self.lines == None:
                self.lines = lines
            if 0 <= line_index < len(lines):
                return lines[line_index].strip()
            else:
                raise IndexError("Line index out of range.")

    def len_lines(self):
        with open(self.filename, "r", encoding="utf-8") as file:
            lines = file.readlines()
            lines = [line for line in lines if line.strip()]  # Ignore empty lines
        return len(lines)

    def append_line(self, line):
        """Append a line to the end of the file."""
        with open(self.filename, "a") as file:
            file.write(line + "\n")

    def load_lines(self):
        """Load all lines from the file into a list."""
        with open(self.filename, "r", encoding="utf-8") as file:
            return [line.strip() for line in file.readlines()]

    def escape_quotes(self, text):
        return text.replace('"', '\\"').replace("'", "\\'")

    def unescape_quotes(self, text):
        return text.replace('\\"', '"').replace("\\'", "'")


def get_safe_path(base_dir, filename):
    """Sanitize and resolve filename, preventing path traversal attacks."""
    safe_name = os.path.basename(filename)
    if not safe_name or safe_name in (".", ".."):
        raise ValueError("Invalid filename")
    base_resolved = os.path.realpath(base_dir)
    resolved = os.path.realpath(os.path.join(base_resolved, safe_name))
    if os.path.commonpath([base_resolved, resolved]) != base_resolved:
        raise ValueError("Path traversal detected")
    return resolved


def base64_to_pil(base64_string):
    header, encoded = base64_string.split(",", 1)  # Remove the data URL header
    image_data = base64.b64decode(encoded)
    return Image.open(io.BytesIO(image_data))


class MyCustomError(Exception):
    def __init__(self, message="Something went wrong"):
        super().__init__(message)


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


any_type = AnyType("*")


class FlexibleOptionalInputType(dict):
    def __init__(self, type):
        self.type = type

    def __getitem__(self, key):
        return (self.type,)

    def __contains__(self, key):
        return True


def time_it(func, *args, **kwargs):
    start_time = time.time()  # Record start time
    result = func(*args, **kwargs)  # Call the passed function with arguments
    end_time = time.time()  # Record end time
    print(f"Execution time: {end_time - start_time:.6f} seconds")
    return result


def clean_text(text):
    # Remove duplicate commas
    text = re.sub(r",+", ",", text)

    # Replace any occurrence of " ," with ","
    text = re.sub(r"\s+,", ",", text)

    # Ensure there's a space after commas followed by a word without space
    text = re.sub(r",(\S)", r", \1", text)

    # Replace multiple spaces with a single space
    text = re.sub(r"\s+", " ", text)

    # Replace any occurrence of ".,", ",." with "."
    text = re.sub(r"\.,|,\.", ".", text)

    # Strip leading and trailing spaces
    return text.strip().replace(" .", ".")


def load_yaml_data(_file_path):
    try:
        # Open the file with UTF-8 encoding
        with open(_file_path, "r", encoding="utf-8") as yaml_file:
            # Load YAML content as Python objects
            _yaml_data = yaml.safe_load(yaml_file)

        # Ensure the data is returned as a list
        if isinstance(_yaml_data, list):
            return _yaml_data
        else:
            raise ValueError("YAML content is not a list of objects.")

    except FileNotFoundError:
        print(f"Error: The file '{_file_path}' was not found.")
    except yaml.YAMLError as e:
        print(f"Error: Invalid YAML format. {e}")
    except Exception as e:
        print(f"An error occurred: {e}")

    return None


def get_yaml_names(_folder_path):
    names = []

    for file_name in os.listdir(_folder_path):
        if file_name.endswith(".yaml") or file_name.endswith(".yml"):
            names.append(file_name)

    return names


possible_names = ["comfyui-itools", "ComfyUI-iTools"]


def check_detect_project_dir():
    paths = folder_paths.folder_names_and_paths["custom_nodes"][0]
    for path in paths:
        for name in possible_names:
            proj_dir = os.path.join(path, name)
            if os.path.exists(proj_dir):
                return proj_dir
    raise FileNotFoundError("No valid project directory found on the device.")


project_dir = check_detect_project_dir()


def get_user_extra_style_choice():
    try:
        ud_dir = os.path.join(folder_paths.base_path, "user", "default")
        settings_file = os.path.join(ud_dir, "comfy.settings.json")
        with open(settings_file, "r") as file:
            settings = json.load(file)
        return settings.get("iTools.Nodes.More Styles", False)
    except (OSError, json.JSONDecodeError, AttributeError):
        return False


allow_extra_styles = get_user_extra_style_choice()

styles = get_yaml_names(os.path.join(project_dir, "styles"))

if allow_extra_styles:
    more_styles = get_yaml_names(os.path.join(project_dir, "styles", "more examples"))
    styles.extend(x for x in more_styles if x not in styles)


def read_styles(_yaml_data):
    if not isinstance(_yaml_data, list):
        print("Error: input data must be a list")
        return None

    names = []

    for item in _yaml_data:
        if isinstance(item, dict):
            if "name" in item:
                names.append(item["name"])

    return names


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


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


def pil2mask(image):
    # Convert grayscale image to numpy array
    numpy_array = torch.tensor(np.array(image), dtype=torch.float32)
    # Normalize to binary mask: 0 for transparent (or dark), 1 for opaque (or bright)
    mask = (numpy_array > 0).float()  # Converts values to 1 if > 0, otherwise 0
    # Adding a batch dimension
    if mask.dim() == 2:
        mask = mask.unsqueeze(0)  # Shape: [1, H, W]
    return mask


# get allow beta nodes
def get_user_dev_mode():
    try:
        ud_dir = os.path.join(folder_paths.base_path, "user", "default")
        settings_file = os.path.join(ud_dir, "comfy.settings.json")
        with open(settings_file, "r") as file:
            settings = json.load(file)
        return settings.get("iTools.Nodes.Dev Mode", True)
    except (OSError, json.JSONDecodeError, AttributeError):
        return True


# get allow dev nodes
def get_user_dev_mode2():
    try:
        ud_dir = os.path.join(folder_paths.base_path, "user", "default")
        settings_file = os.path.join(ud_dir, "comfy.settings.json")
        with open(settings_file, "r") as file:
            settings = json.load(file)
        return settings.get("iTools.Nodes.Dev Mode2", False)
    except (OSError, json.JSONDecodeError, AttributeError):
        return False


# get node display name preferences
def get_user_node_display_name_preferences():
    try:
        ud_dir = os.path.join(folder_paths.base_path, "user", "default")
        settings_file = os.path.join(ud_dir, "comfy.settings.json")
        with open(settings_file, "r") as file:
            settings = json.load(file)
        return settings.get("iTools.Nodes.Node Display Name Preferences", True)
    except (OSError, json.JSONDecodeError, AttributeError):
        return True