import os
import random
import sys
import json
import argparse
import contextlib
from typing import Sequence, Mapping, Any, Union
import torch
from random import randrange


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:
        if args is None or args.comfyui_directory is None:
            path = os.getcwd()
        else:
            path = args.comfyui_directory

    # 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)

        manager_path = os.path.join(
            comfyui_path, "custom_nodes", "ComfyUI-Manager", "glob"
        )

        if os.path.isdir(manager_path) and os.listdir(manager_path):
            sys.path.append(manager_path)
            global has_manager
            has_manager = True

        import __main__

        if getattr(__main__, "__file__", None) is None:
            __main__.__file__ = os.path.join(comfyui_path, "main.py")

        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.
    """
    from comfy.options import enable_args_parsing

    enable_args_parsing()
    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.")


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.
    """
    if has_manager:
        try:
            import manager_core as manager
        except ImportError:
            print("Could not import manager_core, proceeding without it.")
            return
        else:
            if hasattr(manager, "get_config"):
                print("Patching manager_core.get_config to enforce offline mode.")
                try:
                    get_config = manager.get_config

                    def _get_config(*args, **kwargs):
                        config = get_config(*args, **kwargs)
                        config["network_mode"] = "offline"
                        return config

                    manager.get_config = _get_config
                except Exception as e:
                    print("Failed to patch manager_core.get_config:", e)

    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)

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

        # Initializing custom nodes
        await init_extra_nodes(init_custom_nodes=True)

    loop.run_until_complete(inner())


def save_image_wrapper(context, cls):
    if args.output is None:
        return cls

    from PIL import Image, ImageOps, ImageSequence
    from PIL.PngImagePlugin import PngInfo

    import numpy as np

    class WrappedSaveImage(cls):
        counter = 0

        def save_images(
            self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None
        ):
            if args.output is None:
                return super().save_images(
                    images, filename_prefix, prompt, extra_pnginfo
                )
            else:
                if len(images) > 1 and args.output == "-":
                    raise ValueError("Cannot save multiple images to stdout")
                filename_prefix += self.prefix_append

                results = list()
                for batch_number, image in enumerate(images):
                    i = 255.0 * image.cpu().numpy()
                    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
                    metadata = None
                    if not args.disable_metadata:
                        metadata = PngInfo()
                        if prompt is not None:
                            metadata.add_text("prompt", json.dumps(prompt))
                        if extra_pnginfo is not None:
                            for x in extra_pnginfo:
                                metadata.add_text(x, json.dumps(extra_pnginfo[x]))

                    if args.output == "-":
                        # Hack to briefly restore stdout
                        if context is not None:
                            context.__exit__(None, None, None)
                        try:
                            img.save(
                                sys.stdout.buffer,
                                format="png",
                                pnginfo=metadata,
                                compress_level=self.compress_level,
                            )
                        finally:
                            if context is not None:
                                context.__enter__()
                    else:
                        subfolder = ""
                        if len(images) == 1:
                            if os.path.isdir(args.output):
                                subfolder = args.output
                                file = "output.png"
                            else:
                                subfolder, file = os.path.split(args.output)
                                if subfolder == "":
                                    subfolder = os.getcwd()
                        else:
                            if os.path.isdir(args.output):
                                subfolder = args.output
                                file = filename_prefix
                            else:
                                subfolder, file = os.path.split(args.output)

                            if subfolder == "":
                                subfolder = os.getcwd()

                            files = os.listdir(subfolder)
                            file_pattern = file
                            while True:
                                filename_with_batch_num = file_pattern.replace(
                                    "%batch_num%", str(batch_number)
                                )
                                file = (
                                    f"{filename_with_batch_num}_{self.counter:05}.png"
                                )
                                self.counter += 1

                                if file not in files:
                                    break

                        img.save(
                            os.path.join(subfolder, file),
                            pnginfo=metadata,
                            compress_level=self.compress_level,
                        )
                        print("Saved image to", os.path.join(subfolder, file))
                        results.append(
                            {
                                "filename": file,
                                "subfolder": subfolder,
                                "type": self.type,
                            }
                        )

                return {"ui": {"images": results}}

    return WrappedSaveImage


def parse_arg(s: Any, default: Any = None) -> Any:
    """Parses a JSON string, returning it unchanged if the parsing fails."""
    if __name__ == "__main__" or not isinstance(s, str):
        return s

    try:
        return json.loads(s)
    except json.JSONDecodeError:
        return s


parser = argparse.ArgumentParser(
    description="A converted ComfyUI workflow. Node inputs listed below. Values passed should be valid JSON (assumes string if not valid JSON)."
)
parser.add_argument(
    "--ckpt_name1",
    default="SDXLCheckpoint.safetensors",
    help='Argument 0, input `ckpt_name` for node "Load Checkpoint" id 14 (autogenerated)',
)

parser.add_argument(
    "--lora_name2",
    default="dmd2_sdxl_4step_lora_fp16.safetensors",
    help='Argument 2, input `lora_name` for node "Load LoRA" id 17 (autogenerated)',
)

parser.add_argument(
    "--strength_model3",
    default=1,
    help='Argument 3, input `strength_model` for node "Load LoRA" id 17 (autogenerated)',
)

parser.add_argument(
    "--strength_clip4",
    default=1,
    help='Argument 4, input `strength_clip` for node "Load LoRA" id 17 (autogenerated)',
)

parser.add_argument(
    "--text5",
    default="Xx_negative_xX",
    help='Argument 0, input `text` for node "CLIP Text Encode (Prompt)" id 7 (autogenerated)',
)

parser.add_argument(
    "--text6",
    default="Xx_positive_xX",
    help='Argument 0, input `text` for node "CLIPTextEncode with BREAK syntax" id 15 (autogenerated)',
)

parser.add_argument(
    "--image7",
    default="example.png",
    help='Argument 0, input `image` for node "Load Image" id 19 (autogenerated)',
)

parser.add_argument(
    "--block_number8",
    default=3,
    help='Argument 1, input `block_number` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--downscale_factor9",
    default=2,
    help='Argument 2, input `downscale_factor` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--start_percent10",
    default=0,
    help='Argument 3, input `start_percent` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--end_percent11",
    default=0.5000000000000001,
    help='Argument 4, input `end_percent` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--downscale_after_skip12",
    default=True,
    help='Argument 5, input `downscale_after_skip` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--downscale_method13",
    default="bicubic",
    help='Argument 6, input `downscale_method` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--upscale_method14",
    default="bicubic",
    help='Argument 7, input `upscale_method` for node "PatchModelAddDownscale (Kohya Deep Shrink)" id 16 (autogenerated)',
)

parser.add_argument(
    "--seed15",
    default=64836095259134,
    help='Argument 1, input `seed` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--steps16",
    default=8,
    help='Argument 2, input `steps` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--cfg17",
    default=1,
    help='Argument 3, input `cfg` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--sampler_name18",
    default="lcm",
    help='Argument 4, input `sampler_name` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--scheduler19",
    default="karras",
    help='Argument 5, input `scheduler` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--denoise20",
    default=0.87,
    help='Argument 9, input `denoise` for node "KSampler" id 3 (autogenerated)',
)

parser.add_argument(
    "--filename_prefix21",
    default="Fast",
    help='Argument 1, input `filename_prefix` for node "Save Image" id 9 (autogenerated)',
)

parser.add_argument(
    "--queue-size",
    "-q",
    type=int,
    default=1,
    help="How many times the workflow will be executed (default: 1)",
)

parser.add_argument(
    "--comfyui-directory",
    "-c",
    default=None,
    help="Where to look for ComfyUI (default: current directory)",
)

parser.add_argument(
    "--output",
    "-o",
    default=None,
    help="The location to save the output image. Either a file path, a directory, or - for stdout (default: the ComfyUI output directory)",
)

parser.add_argument(
    "--disable-metadata",
    action="store_true",
    help="Disables writing workflow metadata to the outputs",
)


comfy_args = [sys.argv[0]]
if __name__ == "__main__" and "--" in sys.argv:
    idx = sys.argv.index("--")
    comfy_args += sys.argv[idx + 1 :]
    sys.argv = sys.argv[:idx]

args = None
if __name__ == "__main__":
    args = parser.parse_args()
    sys.argv = comfy_args
if args is not None and args.output is not None and args.output == "-":
    ctx = contextlib.redirect_stdout(sys.stderr)
else:
    ctx = contextlib.nullcontext()

PROMPT_DATA = json.loads(
    '{"3": {"inputs": {"seed": 64836095259134, "steps": 8, "cfg": 1, "sampler_name": "lcm", "scheduler": "karras", "denoise": 0.87, "model": ["16", 0], "positive": ["15", 0], "negative": ["7", 0], "latent_image": ["18", 0]}, "class_type": "KSampler", "_meta": {"title": "KSampler"}}, "7": {"inputs": {"text": "Xx_negative_xX", "clip": ["17", 1]}, "class_type": "CLIPTextEncode", "_meta": {"title": "CLIP Text Encode (Prompt)"}}, "8": {"inputs": {"samples": ["3", 0], "vae": ["14", 2]}, "class_type": "VAEDecode", "_meta": {"title": "VAE Decode"}}, "9": {"inputs": {"filename_prefix": "Fast", "images": ["8", 0]}, "class_type": "SaveImage", "_meta": {"title": "Save Image"}}, "14": {"inputs": {"ckpt_name": "SDXLCheckpoint.safetensors"}, "class_type": "CheckpointLoaderSimple", "_meta": {"title": "Load Checkpoint"}}, "15": {"inputs": {"text": "Xx_positive_xX", "clip": ["17", 1]}, "class_type": "CLIPTextEncodeWithBreak", "_meta": {"title": "CLIPTextEncode with BREAK syntax"}}, "16": {"inputs": {"block_number": 3, "downscale_factor": 2, "start_percent": 0, "end_percent": 0.5000000000000001, "downscale_after_skip": true, "downscale_method": "bicubic", "upscale_method": "bicubic", "model": ["17", 0]}, "class_type": "PatchModelAddDownscale", "_meta": {"title": "PatchModelAddDownscale (Kohya Deep Shrink)"}}, "17": {"inputs": {"lora_name": "dmd2_sdxl_4step_lora_fp16.safetensors", "strength_model": 1, "strength_clip": 1, "model": ["14", 0], "clip": ["14", 1]}, "class_type": "LoraLoader", "_meta": {"title": "Load LoRA"}}, "18": {"inputs": {"pixels": ["19", 0], "vae": ["14", 2]}, "class_type": "VAEEncode", "_meta": {"title": "VAE Encode"}}, "19": {"inputs": {"image": "example.png"}, "class_type": "LoadImage", "_meta": {"title": "Load Image"}}}'
)


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.
    """
    if has_manager:
        try:
            import manager_core as manager
        except ImportError:
            print("Could not import manager_core, proceeding without it.")
            return
        else:
            if hasattr(manager, "get_config"):
                print("Patching manager_core.get_config to enforce offline mode.")
                try:
                    get_config = manager.get_config

                    def _get_config(*args, **kwargs):
                        config = get_config(*args, **kwargs)
                        config["network_mode"] = "offline"
                        return config

                    manager.get_config = _get_config
                except Exception as e:
                    print("Failed to patch manager_core.get_config:", e)

    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)

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

        # Initializing custom nodes
        await init_extra_nodes(init_custom_nodes=True)

    loop.run_until_complete(inner())


_custom_nodes_imported = False
_custom_path_added = False


def main(*func_args, **func_kwargs):
    global args, _custom_nodes_imported, _custom_path_added
    if __name__ == "__main__":
        if args is None:
            args = parser.parse_args()
    else:
        defaults = dict(
            (arg, parser.get_default(arg))
            for arg in ["queue_size", "comfyui_directory", "output", "disable_metadata"]
            + [
                "ckpt_name1",
                "lora_name2",
                "strength_model3",
                "strength_clip4",
                "text5",
                "text6",
                "image7",
                "block_number8",
                "downscale_factor9",
                "start_percent10",
                "end_percent11",
                "downscale_after_skip12",
                "downscale_method13",
                "upscale_method14",
                "seed15",
                "steps16",
                "cfg17",
                "sampler_name18",
                "scheduler19",
                "denoise20",
                "filename_prefix21",
            ]
        )

        all_args = dict()
        all_args.update(defaults)
        all_args.update(func_kwargs)

        args = argparse.Namespace(**all_args)

    with ctx:
        if not _custom_path_added:
            add_comfyui_directory_to_sys_path()
            add_extra_model_paths()

            _custom_path_added = True

        if not _custom_nodes_imported:
            import_custom_nodes()

            _custom_nodes_imported = True

        from nodes import NODE_CLASS_MAPPINGS

    with torch.inference_mode(), ctx:
        checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
        checkpointloadersimple_14 = checkpointloadersimple.load_checkpoint(
            ckpt_name=parse_arg(args.ckpt_name1)
        )

        loraloader = NODE_CLASS_MAPPINGS["LoraLoader"]()
        loraloader_17 = loraloader.load_lora(
            lora_name=parse_arg(args.lora_name2),
            strength_model=parse_arg(args.strength_model3),
            strength_clip=parse_arg(args.strength_clip4),
            model=get_value_at_index(checkpointloadersimple_14, 0),
            clip=get_value_at_index(checkpointloadersimple_14, 1),
        )

        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
        cliptextencode_7 = cliptextencode.encode(
            text=parse_arg(args.text5), clip=get_value_at_index(loraloader_17, 1)
        )

        cliptextencodewithbreak = NODE_CLASS_MAPPINGS["CLIPTextEncodeWithBreak"]()
        cliptextencodewithbreak_15 = cliptextencodewithbreak.encode(
            text=parse_arg(args.text6), clip=get_value_at_index(loraloader_17, 1)
        )

        loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
        loadimage_19 = loadimage.load_image(image=parse_arg(args.image7))

        vaeencode = NODE_CLASS_MAPPINGS["VAEEncode"]()
        vaeencode_18 = vaeencode.encode(
            pixels=get_value_at_index(loadimage_19, 0),
            vae=get_value_at_index(checkpointloadersimple_14, 2),
        )

        patchmodeladddownscale = NODE_CLASS_MAPPINGS["PatchModelAddDownscale"]()
        ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
        vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
        saveimage = save_image_wrapper(ctx, NODE_CLASS_MAPPINGS["SaveImage"])()
        for q in range(args.queue_size):
            patchmodeladddownscale_16 = patchmodeladddownscale.patch(
                block_number=parse_arg(args.block_number8),
                downscale_factor=parse_arg(args.downscale_factor9),
                start_percent=parse_arg(args.start_percent10),
                end_percent=parse_arg(args.end_percent11),
                downscale_after_skip=parse_arg(args.downscale_after_skip12),
                downscale_method=parse_arg(args.downscale_method13),
                upscale_method=parse_arg(args.upscale_method14),
                model=get_value_at_index(loraloader_17, 0),
            )

            ksampler_3 = ksampler.sample(
                seed=randrange(1000000000),
                steps=parse_arg(args.steps16),
                cfg=parse_arg(args.cfg17),
                sampler_name=parse_arg(args.sampler_name18),
                scheduler=parse_arg(args.scheduler19),
                denoise=(90.0+randrange(9))/100.0,
                model=get_value_at_index(patchmodeladddownscale_16, 0),
                positive=get_value_at_index(cliptextencodewithbreak_15, 0),
                negative=get_value_at_index(cliptextencode_7, 0),
                latent_image=get_value_at_index(vaeencode_18, 0),
            )

            vaedecode_8 = vaedecode.decode(
                samples=get_value_at_index(ksampler_3, 0),
                vae=get_value_at_index(checkpointloadersimple_14, 2),
            )

            if __name__ != "__main__":
                return dict(
                    filename_prefix=parse_arg(args.filename_prefix21),
                    images=get_value_at_index(vaedecode_8, 0),
                    prompt=PROMPT_DATA,
                )
            else:
                saveimage_9 = saveimage.save_images(
                    filename_prefix=parse_arg(args.filename_prefix21),
                    images=get_value_at_index(vaedecode_8, 0),
                    prompt=PROMPT_DATA,
                )


if __name__ == "__main__":
    main()