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")
hf_hub_download(repo_id="ximso/RealESRGAN_x4plus_anime_6B", filename="RealESRGAN_x4plus_anime_6B.pth", local_dir="models/upscale_models")

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
    asyncio.run(init_extra_nodes())


from nodes import NODE_CLASS_MAPPINGS
from comfy_extras.nodes_upscale_model import UpscaleModelLoader

import_custom_nodes()

checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
    ckpt_name="zukiCuteILL_v60.safetensors"
)
cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()

loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
vaeencode = NODE_CLASS_MAPPINGS["VAEEncode"]()
conditioningconcat = NODE_CLASS_MAPPINGS["ConditioningConcat"]()
repeatlatentbatch = NODE_CLASS_MAPPINGS["RepeatLatentBatch"]()
ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
upscalemodelloader_220 = UpscaleModelLoader.execute(
    model_name="RealESRGAN_x4plus_anime_6B.pth"
)
pixelksampleupscalerprovider = NODE_CLASS_MAPPINGS["PixelKSampleUpscalerProvider"]()
iterativelatentupscale = NODE_CLASS_MAPPINGS["IterativeLatentUpscale"]()
stepsschedulehookprovider = NODE_CLASS_MAPPINGS["StepsScheduleHookProvider"]()
cfgschedulehookprovider = NODE_CLASS_MAPPINGS["CfgScheduleHookProvider"]()
pixelksamplehookcombine = NODE_CLASS_MAPPINGS["PixelKSampleHookCombine"]()

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),
)
cfgschedulehookprovider_541 = cfgschedulehookprovider.doit(
    schedule_for_iteration="simple", target_cfg=10
)

@spaces.GPU
def generate_image(param_image, param_prompt, param_creative, param_style, param_prefix):

    param_creative = float(param_creative)

    if param_creative > 0.35:
        param_amount1 = 3
        param_amount2 = 1
        param_step = 7
        param_step2 = 15
    else:
        param_amount1 = 1
        param_amount2 = 3
        param_step = 8
        param_step2 = 17

    with torch.inference_mode():

        loadimage_89 = loadimage.load_image(image=param_image)
        vaeencode_229 = vaeencode.encode(
            pixels=get_value_at_index(loadimage_89, 0),
            vae=get_value_at_index(checkpointloadersimple_4, 2),
        )
        cliptextencode_524 = cliptextencode.encode(
            text=param_prompt,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        cliptextencode_526 = cliptextencode.encode(
            text=param_style,
            clip=get_value_at_index(checkpointloadersimple_4, 1),
        )
        conditioningconcat_521 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(cliptextencode_526, 0),
            conditioning_from=get_value_at_index(cliptextencode_524, 0),
        )
        conditioningconcat_527 = conditioningconcat.concat(
            conditioning_to=get_value_at_index(conditioningconcat_521, 0),
            conditioning_from=get_value_at_index(cliptextencode_525, 0),
        )
        repeatlatentbatch_506 = repeatlatentbatch.repeat(
            amount=param_amount1, samples=get_value_at_index(vaeencode_229, 0)
        )
        ksampler_230 = ksampler.sample(
            seed=random.randint(1, 2**64),
            steps=20,
            cfg=6,
            sampler_name="euler_ancestral",
            scheduler="normal",
            denoise=param_creative,
            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(repeatlatentbatch_506, 0),
        )
        repeatlatentbatch_509 = repeatlatentbatch.repeat(
            amount=param_amount2, samples=get_value_at_index(ksampler_230, 0)
        )
        stepsschedulehookprovider_537 = stepsschedulehookprovider.doit(
            schedule_for_iteration="simple", target_steps=param_step2
        )
        pixelksamplehookcombine_540 = pixelksamplehookcombine.doit(
            hook1=get_value_at_index(stepsschedulehookprovider_537, 0),
            hook2=get_value_at_index(cfgschedulehookprovider_541, 0),
        )
        pixelksampleupscalerprovider_462 = pixelksampleupscalerprovider.doit(
            scale_method="lanczos",
            seed=random.randint(1, 2**64),
            steps=param_step,
            cfg=9,
            sampler_name="euler",
            scheduler="normal",
            denoise=0.35,
            use_tiled_vae=False,
            tile_size=512,
            model=get_value_at_index(checkpointloadersimple_4, 0),
            vae=get_value_at_index(checkpointloadersimple_4, 2),
            positive=get_value_at_index(conditioningconcat_527, 0),
            negative=get_value_at_index(cliptextencode_7, 0),
            upscale_model_opt=get_value_at_index(upscalemodelloader_220, 0),
            pk_hook_opt=get_value_at_index(pixelksamplehookcombine_540, 0),
        )
        iterativelatentupscale_461 = iterativelatentupscale.doit(
            upscale_factor=1.5,
            steps=2,
            temp_prefix="",
            step_mode="simple",
            samples=get_value_at_index(repeatlatentbatch_509, 0),
            upscaler=get_value_at_index(pixelksampleupscalerprovider_462, 0),
            unique_id=1445395014345641493,
        )
        vaedecode_233 = vaedecode.decode(
            samples=get_value_at_index(iterativelatentupscale_461, 0),
            vae=get_value_at_index(iterativelatentupscale_461, 1),
        )
        saveimage_410 = saveimage.save_images(
            filename_prefix=param_prefix,
            images=get_value_at_index(vaedecode_233, 0),
        )
        saved_path = [
          f"output/{saveimage_410['ui']['images'][0]['filename']}",
          f"output/{saveimage_410['ui']['images'][1]['filename']}",
          f"output/{saveimage_410['ui']['images'][2]['filename']}",
        ]
        return saved_path

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(label="Image", type="filepath", height=300, show_label=False)
            prompt = gr.Textbox(label="prompt", lines=3, max_lines=3, placeholder="prompt")
            style = gr.Textbox(label="style", lines=2, max_lines=2, placeholder="style")
            creative = gr.Dropdown(
                choices=[
                    ("balance", 0.65),
                    ("none", 0),
                    ("low", 0.25),
                    ("normal", 0.5),
                    ("high", 0.75),
                    ("ultra", 1),
                ],
                allow_custom_value=True,
                value=0.65,
                label="creative"
            )
            run_btn = gr.Button("Generate", variant="primary")
            prefix = gr.Textbox(visible=False, value="comfyui_")
        with gr.Column(scale=2):
            output_image = gr.Gallery(
                label="Result",
                columns=3,
                object_fit="contain",
                height="auto"
            )
    run_btn.click(
        fn=generate_image,
        inputs=[image, prompt, creative, style, prefix],
        outputs=[output_image]
    )

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