import os, glob, random
from collections import Counter
from PIL import Image
from math import isqrt, ceil
from typing import List
import logging
import hashlib
import torch
from safetensors.torch import save_file, safe_open
from insightface.app.common import Face

from modules.images import FilenameGenerator, get_next_sequence_number
from modules import shared, script_callbacks
from scripts.reactor_globals import DEVICE, BASE_PATH, FACE_MODELS_PATH, IS_SDNEXT

try:
    from modules.paths_internal import models_path
except:
    try:
        from modules.paths import models_path
    except:
        model_path = os.path.abspath("models")

MODELS_PATH = None

def set_Device(value):
    global DEVICE
    DEVICE = value
    with open(os.path.join(BASE_PATH, "last_device.txt"), "w") as txt:
        txt.write(DEVICE)

def get_Device():
    global DEVICE
    return DEVICE

def set_SDNEXT():
    global IS_SDNEXT
    IS_SDNEXT = True

def get_SDNEXT():
    global IS_SDNEXT
    return IS_SDNEXT

def make_grid(image_list: List):
    
    # Count the occurrences of each image size in the image_list
    size_counter = Counter(image.size for image in image_list)
    
    # Get the most common image size (size with the highest count)
    common_size = size_counter.most_common(1)[0][0]
    
    # Filter the image_list to include only images with the common size
    image_list = [image for image in image_list if image.size == common_size]
    
    # Get the dimensions (width and height) of the common size
    size = common_size
    
    # If there are more than one image in the image_list
    if len(image_list) > 1:
        num_images = len(image_list)
        
        # Calculate the number of rows and columns for the grid
        rows = isqrt(num_images)
        cols = ceil(num_images / rows)

        # Calculate the size of the square image
        square_size = (cols * size[0], rows * size[1])

        # Create a new RGB image with the square size
        square_image = Image.new("RGB", square_size)

        # Paste each image onto the square image at the appropriate position
        for i, image in enumerate(image_list):
            row = i // cols
            col = i % cols

            square_image.paste(image, (col * size[0], row * size[1]))

        # Return the resulting square image
        return square_image
    
    # Return None if there are no images or only one image in the image_list
    return None

def get_image_path(image, path, basename, seed=None, prompt=None, extension='png', p=None, suffix=""):
    
    namegen = FilenameGenerator(p, seed, prompt, image)

    save_to_dirs = shared.opts.save_to_dirs

    if save_to_dirs:
        dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /')
        path = os.path.join(path, dirname)

    os.makedirs(path, exist_ok=True)

    if seed is None:
        file_decoration = ""
    elif shared.opts.save_to_dirs:
        file_decoration = shared.opts.samples_filename_pattern or "[seed]"
    else:
        file_decoration = shared.opts.samples_filename_pattern or "[seed]-[prompt_spaces]"

    file_decoration = namegen.apply(file_decoration) + suffix

    add_number = shared.opts.save_images_add_number or file_decoration == ''

    if file_decoration != "" and add_number:
        file_decoration = f"-{file_decoration}"

    if add_number:
        basecount = get_next_sequence_number(path, basename)
        fullfn = None
        for i in range(500):
            fn = f"{basecount + i:05}" if basename == '' else f"{basename}-{basecount + i:04}"
            fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
            if not os.path.exists(fullfn):
                break
    else:
        fullfn = os.path.join(path, f"{file_decoration}.{extension}")

    pnginfo = {}

    params = script_callbacks.ImageSaveParams(image, p, fullfn, pnginfo)
    # script_callbacks.before_image_saved_callback(params)

    fullfn = params.filename

    fullfn_without_extension, extension = os.path.splitext(params.filename)
    if hasattr(os, 'statvfs'):
        max_name_len = os.statvfs(path).f_namemax
        fullfn_without_extension = fullfn_without_extension[:max_name_len - max(4, len(extension))]
        params.filename = fullfn_without_extension + extension
        fullfn = params.filename

    return fullfn

def addLoggingLevel(levelName, levelNum, methodName=None):
    if not methodName:
        methodName = levelName.lower()

    def logForLevel(self, message, *args, **kwargs):
        if self.isEnabledFor(levelNum):
            self._log(levelNum, message, args, **kwargs)

    def logToRoot(message, *args, **kwargs):
        logging.log(levelNum, message, *args, **kwargs)

    logging.addLevelName(levelNum, levelName)
    setattr(logging, levelName, levelNum)
    setattr(logging.getLoggerClass(), methodName, logForLevel)
    setattr(logging, methodName, logToRoot)

def get_image_md5hash(image: Image.Image):
    md5hash = hashlib.md5(image.tobytes())
    return md5hash.hexdigest()

def save_face_model(face: Face, filename: str) -> None:
    try:
        tensors = {
            "bbox": torch.tensor(face["bbox"]),
            "kps": torch.tensor(face["kps"]),
            "det_score": torch.tensor(face["det_score"]),
            "landmark_3d_68": torch.tensor(face["landmark_3d_68"]),
            "pose": torch.tensor(face["pose"]),
            "landmark_2d_106": torch.tensor(face["landmark_2d_106"]),
            "embedding": torch.tensor(face["embedding"]),
            "gender": torch.tensor(face["gender"]),
            "age": torch.tensor(face["age"]),
        }
        save_file(tensors, filename)
        # print(f"Face model has been saved to '{filename}'")
    except Exception as e:
        print(f"Error: {e}")

def get_models():
    global MODELS_PATH
    models_path_init = os.path.join(models_path, "insightface/*")
    models = glob.glob(models_path_init)
    models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")]
    models_names = []
    for model in models:
        model_path = os.path.split(model)
        if MODELS_PATH is None:
            MODELS_PATH = model_path[0]
        model_name = model_path[1]
        models_names.append(model_name)
    return models_names

def load_face_model(filename: str):
    face = {}
    model_path = os.path.join(FACE_MODELS_PATH, filename)
    with safe_open(model_path, framework="pt") as f:
        for k in f.keys():
            face[k] = f.get_tensor(k).numpy()
    return Face(face)

def get_facemodels():
    models_path = os.path.join(FACE_MODELS_PATH, "*")
    models = glob.glob(models_path)
    models = [x for x in models if x.endswith(".safetensors")]
    return models

def get_model_names(get_models):
    models = get_models()
    names = ["None"]
    for x in models:
        names.append(os.path.basename(x))
    return names

def get_images_from_folder(path: str):
    files_path = os.path.join(path, "*")
    files = glob.glob(files_path)
    images = []
    images_names = []
    for x in files:
        if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp')):
            images.append(Image.open(x))
            images_names.append(os.path.basename(x))
    return images,images_names
    # return [Image.open(x) for x in images if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp'))],[os.path.basename(x) for x in images if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp'))]

def get_random_image_from_folder(path: str):
    images,names = get_images_from_folder(path)
    random_image_index = random.randint(0, len(images) - 1)
    return [images[random_image_index]],[names[random_image_index]]

def get_images_from_list(imgs: List):
    images = []
    images_names = []
    for x in imgs:
        images.append(Image.open(os.path.abspath(x.name)))
        images_names.append(os.path.basename(x.name))
    return images,images_names
    # return [Image.open(os.path.abspath(x.name)) for x in imgs],[os.path.basename(x.name) for x in imgs]
