import torch
import latent_preview

from comfy.cli_args import args
from comfy.samplers import sample

WARNING_COLOR = "\033[93m"  # Yellow


# Extract global preview_method
def global_preview_method():
    return args.preview_method


def format_message(text, color_code):
    RESET_COLOR = "\033[0m"
    return f"{color_code}{text}{RESET_COLOR}"


def warning(text):
    return format_message(text, WARNING_COLOR)


# Set global preview_method
def set_preview_method(method):
    if method == 'auto' or method == 'LatentPreviewMethod.Auto':
        args.preview_method = latent_preview.LatentPreviewMethod.Auto
    elif method == 'latent2rgb' or method == 'LatentPreviewMethod.Latent2RGB':
        args.preview_method = latent_preview.LatentPreviewMethod.Latent2RGB
    elif method == 'taesd' or method == 'LatentPreviewMethod.TAESD':
        args.preview_method = latent_preview.LatentPreviewMethod.TAESD
    else:
        args.preview_method = latent_preview.LatentPreviewMethod.NoPreviews


# Cache for Efficiency Node models
loaded_objects = {
    "ckpt": [], # (ckpt_name, ckpt_model, clip, bvae, [id])
    "refn": [], # (ckpt_name, ckpt_model, clip, bvae, [id])
    "vae": [],  # (vae_name, vae, [id])
    "lora": []  # ([(lora_name, strength_model, strength_clip)], ckpt_name, lora_model, clip_lora, [id])
}


# Cache for Efficient Ksamplers
last_helds = {
    "latent": [],   # (latent, [parameters], id)    # Base sampling latent results
    "image": [],    # (image, id)                   # Base sampling image results
    "cnet_img": []  # (cnet_img, [parameters], id)  # HiRes-Fix control net preprocessor image results
}


# This function cleans global variables associated with nodes that are no longer detected on UI
def globals_cleanup(prompt):
    global loaded_objects
    global last_helds

    # Step 1: Clean up last_helds
    for key in list(last_helds.keys()):
        original_length = len(last_helds[key])
        last_helds[key] = [
            (*values, id_)
            for *values, id_ in last_helds[key]
            if str(id_) in prompt.keys()
        ]

    # Step 2: Clean up loaded_objects
    for key in list(loaded_objects.keys()):
        for i, tup in enumerate(list(loaded_objects[key])):
            # Remove ids from id array in each tuple that don't exist in prompt
            id_array = [id for id in tup[-1] if str(id) in prompt.keys()]
            if len(id_array) != len(tup[-1]):
                if id_array:
                    loaded_objects[key][i] = tup[:-1] + (id_array,)
                    #print(f'Updated tuple at index {i} in {key} in loaded_objects: {loaded_objects[key][i]}')
                else:
                    # If id array becomes empty, delete the corresponding tuple
                    loaded_objects[key].remove(tup)
                    #print(f'Deleted tuple at index {i} in {key} in loaded_objects because its id array became empty.')


def sample_custom_ultra(model, device, noise, sampler, positive, negative, cfg, model_options={}, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):    
    if last_step is not None and last_step < (len(sigmas) - 1):
        sigmas = sigmas[:last_step + 1]
        if force_full_denoise:
            sigmas[-1] = 0

    if start_step is not None:
        if start_step < (len(sigmas) - 1):
            sigmas = sigmas[start_step:]
        else:
            if latent_image is not None:
                return latent_image
            else:
                return torch.zeros_like(noise)

    return sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)


def sample_custom_ultra_plus(noise, sampler, guider, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
    if last_step is not None and last_step < (len(sigmas) - 1):
        sigmas = sigmas[:last_step + 1]
        if force_full_denoise:
            sigmas[-1] = 0

    if start_step is not None:
        if start_step < (len(sigmas) - 1):
            sigmas = sigmas[start_step:]
        else:
            if latent_image is not None:
                return latent_image
            else:
                return torch.zeros_like(noise)
    
    return guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)