import spaces
import os
import random
import sys
from typing import Sequence, Mapping, Any, Union
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from comfy import model_management

hf_hub_download(repo_id="John6666/zuki-cute-ill-v60-sdxl", filename="zukiCuteILL_v60.safetensors", local_dir="models/checkpoints")

def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
    """Returns the value at the given index of a sequence or mapping.

    If the object is a sequence (like list or string), returns the value at the given index.
    If the object is a mapping (like a dictionary), returns the value at the index-th key.

    Some return a dictionary, in these cases, we look for the "results" key

    Args:
        obj (Union[Sequence, Mapping]): The object to retrieve the value from.
        index (int): The index of the value to retrieve.

    Returns:
        Any: The value at the given index.

    Raises:
        IndexError: If the index is out of bounds for the object and the object is not a mapping.
    """
    try:
        return obj[index]
    except KeyError:
        return obj["result"][index]


def find_path(name: str, path: str = None) -> str:
    """
    Recursively looks at parent folders starting from the given path until it finds the given name.
    Returns the path as a Path object if found, or None otherwise.
    """
    # If no path is given, use the current working directory
    if path is None:
        path = os.getcwd()

    # Check if the current directory contains the name
    if name in os.listdir(path):
        path_name = os.path.join(path, name)
        print(f"{name} found: {path_name}")
        return path_name

    # Get the parent directory
    parent_directory = os.path.dirname(path)

    # If the parent directory is the same as the current directory, we've reached the root and stop the search
    if parent_directory == path:
        return None

    # Recursively call the function with the parent directory
    return find_path(name, parent_directory)


def add_comfyui_directory_to_sys_path() -> None:
    """
    Add 'ComfyUI' to the sys.path
    """
    comfyui_path = find_path("ComfyUI")
    if comfyui_path is not None and os.path.isdir(comfyui_path):
        sys.path.append(comfyui_path)
        print(f"'{comfyui_path}' added to sys.path")


def add_extra_model_paths() -> None:
    """
    Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
    """
    try:
        from main import load_extra_path_config
    except ImportError:
        print(
            "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
        )
        from utils.extra_config import load_extra_path_config

    extra_model_paths = find_path("extra_model_paths.yaml")

    if extra_model_paths is not None:
        load_extra_path_config(extra_model_paths)
    else:
        print("Could not find the extra_model_paths config file.")


add_comfyui_directory_to_sys_path()
add_extra_model_paths()


def import_custom_nodes() -> None:
    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS

    This function sets up a new asyncio event loop, initializes the PromptServer,
    creates a PromptQueue, and initializes the custom nodes.
    """
    import asyncio
    import execution
    from nodes import init_extra_nodes
    import server

    # Creating a new event loop and setting it as the default loop
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    # Creating an instance of PromptServer with the loop
    server_instance = server.PromptServer(loop)
    execution.PromptQueue(server_instance)

    # Initializing custom nodes
    init_extra_nodes()


from nodes import NODE_CLASS_MAPPINGS

import_custom_nodes()

checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
    ckpt_name="zukiCuteILL_v60.safetensors"
)
cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
conditioningconcat = NODE_CLASS_MAPPINGS["ConditioningConcat"]()
ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()

model_loaders = [checkpointloadersimple_4]

# 检查哪些模型是有效的，并确定最佳加载方式
valid_models = [
    getattr(loader[0], 'patcher', loader[0])
    for loader in model_loaders
    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
]

# 最终加载模型
model_management.load_models_gpu(valid_models)

cliptextencode_7 = cliptextencode.encode(
    text="lowres, bad quality, worst quality, bad anatomy, sketch, jpeg artifacts, ugly, poorly drawn, (signature, watermark, username, logo, web address, twitter_username, patreon_username, character_name, copyright_name), (censored, mosaic_censoring, convenient_censoring, bar_censor, heart_censor), blurry, simple background, transparent background,",
    clip=get_value_at_index(checkpointloadersimple_4, 1),
)
cliptextencode_525 = cliptextencode.encode(
    text="masterpiece, best quality, amazing quality, very aesthetic, absurdres, newest, volumetric lighting, dramatic lighting, ",
    clip=get_value_at_index(checkpointloadersimple_4, 1),
)

@spaces.GPU
def generate_image(param_prompt, param_s1, param_s2, param_s3, param_size_str, param_seed, param_prefix):
    param_width, param_height = map(int, param_size_str.split("x"))
    param_actual_seed1 = param_seed if param_seed != -1 else random.randint(1, 2**64)
    param_actual_seed2 = param_seed if param_seed != -1 else random.randint(1, 2**64)
    param_actual_seed3 = param_seed if param_seed != -1 else random.randint(1, 2**64)

    with torch.inference_mode():

        cliptextencode_524 = cliptextencode.encode(
            text=param_prompt,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        cliptextencode_526 = cliptextencode.encode(
            text=param_s1,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        emptylatentimage_529 = emptylatentimage.generate(
            width=param_width, height=param_height, batch_size=1
        )
        emptylatentimage_530 = emptylatentimage.generate(
            width=param_width, height=param_height, batch_size=1
        )
        emptylatentimage_531 = emptylatentimage.generate(
            width=param_width, height=param_height, batch_size=1
        )
        cliptextencode_540 = cliptextencode.encode(
            text=param_s2,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        cliptextencode_541 = cliptextencode.encode(
            text=param_s3,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        conditioningconcat_521 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(cliptextencode_525, 0),
            conditioning_from=get_value_at_index(cliptextencode_524, 0),
        )
        conditioningconcat_527 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(cliptextencode_526, 0),
            conditioning_from=get_value_at_index(conditioningconcat_521, 0),
        )
        ksampler_230 = ksampler.sample(
            seed=param_actual_seed1,
            steps=25,
            cfg=8,
            sampler_name="euler_ancestral",
            scheduler="normal",
            denoise=1,
            model=get_value_at_index(checkpointloadersimple_4, 0),
            positive=get_value_at_index(conditioningconcat_527, 0),
            negative=get_value_at_index(cliptextencode_7, 0),
            latent_image=get_value_at_index(emptylatentimage_531, 0),
        )
        vaedecode_233 = vaedecode.decode(
            samples=get_value_at_index(ksampler_230, 0),
            vae=get_value_at_index(checkpointloadersimple_4, 2),
        )
        saveimage_410 = saveimage.save_images(
            filename_prefix=param_prefix,
            images=get_value_at_index(vaedecode_233, 0),
        )
        conditioningconcat_543 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(cliptextencode_540, 0),
            conditioning_from=get_value_at_index(conditioningconcat_521, 0),
        )
        ksampler_532 = ksampler.sample(
            seed=param_actual_seed2,
            steps=25,
            cfg=8,
            sampler_name="euler_ancestral",
            scheduler="normal",
            denoise=1,
            model=get_value_at_index(checkpointloadersimple_4, 0),
            positive=get_value_at_index(conditioningconcat_543, 0),
            negative=get_value_at_index(cliptextencode_7, 0),
            latent_image=get_value_at_index(emptylatentimage_530, 0),
        )
        conditioningconcat_544 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(cliptextencode_541, 0),
            conditioning_from=get_value_at_index(conditioningconcat_521, 0),
        )
        ksampler_533 = ksampler.sample(
            seed=param_actual_seed3,
            steps=25,
            cfg=8,
            sampler_name="euler_ancestral",
            scheduler="normal",
            denoise=1,
            model=get_value_at_index(checkpointloadersimple_4, 0),
            positive=get_value_at_index(conditioningconcat_544, 0),
            negative=get_value_at_index(cliptextencode_7, 0),
            latent_image=get_value_at_index(emptylatentimage_529, 0),
        )
        vaedecode_535 = vaedecode.decode(
            samples=get_value_at_index(ksampler_532, 0),
            vae=get_value_at_index(checkpointloadersimple_4, 2),
        )
        saveimage_534 = saveimage.save_images(
            filename_prefix=param_prefix,
            images=get_value_at_index(vaedecode_535, 0),
        )
        vaedecode_537 = vaedecode.decode(
            samples=get_value_at_index(ksampler_533, 0),
            vae=get_value_at_index(checkpointloadersimple_4, 2),
        )
        saveimage_536 = saveimage.save_images(
            filename_prefix=param_prefix,
            images=get_value_at_index(vaedecode_537, 0),
        )
        saved_path = [
          f"output/{saveimage_410['ui']['images'][0]['filename']}",
          f"output/{saveimage_534['ui']['images'][0]['filename']}",
          f"output/{saveimage_536['ui']['images'][0]['filename']}",
        ]
        return saved_path

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="")
            style1 = gr.Textbox(show_label=False, lines=3, placeholder="style1")
            style2 = gr.Textbox(show_label=False, lines=3, placeholder="style2")
            style3 = gr.Textbox(show_label=False, lines=3, placeholder="style3")
            size = gr.Radio(
                show_label=False,
                choices=[
                    ("vertical", "768x1152"),
                    ("horizontal", "1152x768"),
                    ("square", "960x960")
                ],
                value="768x1152",
            )
            seed_input = gr.Number(label="seed", value=-1, precision=0)
            prefix = gr.Textbox(visible=False, value="comfyui_")
            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=2):
            output_image = gr.Gallery(
                label="Result",
                columns=3,
                rows=1,
                object_fit="contain"
            )
    run_btn.click(
        fn=generate_image,
        inputs=[prompt, style1, style2, style3, size, seed_input, prefix],
        outputs=[output_image]
    )

if __name__ == "__main__":
    app.launch(share=True)