# noqa
import json
import logging
import sys
import traceback
import os
from typing import Union

from serverless_comfy.utils import timed
from serverless_comfy.utils import get_comfyui_version


def patch_overrides():
    patch_validate_prompt()
    patch_recursive_execute()
    patch_gligen_set_position()
    patch_init_extra_nodes()
    patch_send_json()
    patch_unet16_fp16_to_bf16()


def patch_validate_prompt():
    from main import execution

    if get_comfyui_version(os.getcwd()) >= (0, 3, 44):
        execution.validate_prompt = avalidate_prompt
    else:
        execution.validate_prompt = validate_prompt


def patch_recursive_execute():
    from main import execution

    execution.recursive_execute = recursive_execute


def patch_gligen_set_position():
    from comfy.gligen import Gligen

    Gligen.set_position = _gligen_set_position


def patch_init_extra_nodes():
    import nodes

    if get_comfyui_version(os.getcwd()) >= (0, 3, 47):
        nodes.init_extra_nodes = ainit_extra_nodes
    else:
        nodes.init_extra_nodes = init_extra_nodes


def patch_send_json():
    from server import PromptServer

    PromptServer.send_json = send_json


def patch_unet16_fp16_to_bf16():
    from comfy.cli_args import args

    if get_comfyui_version(os.getcwd()) >= (0, 3, 14):
        args.fp16_unet = False


def get_outputs(prompt, partial_execution_list=None):
    import nodes
    from execution import full_type_name

    outputs = set()
    node_imports_validated = {}
    errors = []
    node_import_errors = {}
    for x in prompt:
        try:
            class_type = prompt[x]["class_type"]
            class_ = nodes.NODE_CLASS_MAPPINGS[class_type]
            if hasattr(class_, "OUTPUT_NODE") and class_.OUTPUT_NODE == True:  # noqa: E712
                if partial_execution_list is None or x in partial_execution_list:
                    outputs.add(x)
        except Exception as e:
            typ, _, tb = sys.exc_info()
            exception_type = full_type_name(typ)
            reasons = [
                {
                    "type": "exception_during_validation",
                    "message": (
                        f"Node class {prompt[x]['class_type']} not found. "
                        "Make sure the related custom node is installed."
                    ),
                    "details": str(e),
                    "extra_info": {
                        "exception_type": exception_type,
                        "traceback": traceback.format_tb(tb),
                    },
                }
            ]
            logging.error(f"Failed to validate prompt for output {x}:")
            logging.error("* (prompt):")
            for reason in reasons:
                logging.error(f"  - {reason['message']}: {reason['details']}")

            node_imports_validated[x] = (False, reasons, x)
            errors += [(x, reasons)]
            node_import_errors[x] = {
                "errors": reasons,
                "dependent_outputs": [],
                "class_type": class_type,
            }
            logging.error(f"* {class_type} {x}:")
            for reason in reasons:
                logging.error(f"  - {reason['message']}: {reason['details']}")

    if len(node_import_errors) != 0:
        errors_list = []
        for o, errors in errors:
            for error in errors:
                errors_list.append(f"{error['message']}: {error['details']}")
        errors_list = "\n".join(errors_list)

        error = {
            "type": "prompt_outputs_failed_validation",
            "message": "Prompt outputs failed validation",
            "details": errors_list,
            "extra_info": {},
        }

        return (False, error, [], node_import_errors)

    return (True, outputs)


async def avalidate_prompt(*args):
    from execution import full_type_name, validate_inputs

    comfyui_version = get_comfyui_version(os.getcwd())
    if comfyui_version >= (0, 3, 47):
        prompt_id, prompt, partial_execution_list = args
        partial_execution_list: Union[list[str], None]
    elif comfyui_version >= (0, 3, 44):
        prompt_id, prompt = args
        partial_execution_list = None
    else:
        raise RuntimeError("Incompatible ComfyUI version")

    result = get_outputs(prompt, partial_execution_list)
    if result[0] is not True:
        return result

    outputs = result[1]

    good_outputs = set()
    errors = []
    node_errors = {}
    validated = {}
    for o in outputs:
        valid = False
        reasons = []
        try:
            m = await validate_inputs(prompt_id, prompt, o, validated)
            valid = m[0]
            reasons = m[1]
        except Exception as ex:
            typ, _, tb = sys.exc_info()
            valid = False
            exception_type = full_type_name(typ)
            reasons = [
                {
                    "type": "exception_during_validation",
                    "message": "Exception when validating node",
                    "details": str(ex),
                    "extra_info": {
                        "exception_type": exception_type,
                        "traceback": traceback.format_tb(tb),
                    },
                }
            ]
            validated[o] = (False, reasons, o)

        if valid is True:
            good_outputs.add(o)
        else:
            logging.error(f"Failed to validate prompt for output {o}:")
            if len(reasons) > 0:
                logging.error("* (prompt):")
                for reason in reasons:
                    logging.error(f"  - {reason['message']}: {reason['details']}")
            errors += [(o, reasons)]
            for node_id, result in validated.items():
                valid = result[0]
                reasons = result[1]
                # If a node upstream has errors, the nodes downstream will also
                # be reported as invalid, but there will be no errors attached.
                # So don't return those nodes as having errors in the response.
                if valid is not True and len(reasons) > 0:
                    if node_id not in node_errors:
                        class_type = prompt[node_id]["class_type"]
                        node_errors[node_id] = {
                            "errors": reasons,
                            "dependent_outputs": [],
                            "class_type": class_type,
                        }
                        logging.error(f"* {class_type} {node_id}:")
                        for reason in reasons:
                            logging.error(
                                f"  - {reason['message']}: {reason['details']}"
                            )
                    node_errors[node_id]["dependent_outputs"].append(o)
            logging.error("Output will be ignored")

    if len(good_outputs) == 0:
        errors_list = []
        for o, errors in errors:
            for error in errors:
                errors_list.append(f"{error['message']}: {error['details']}")
        errors_list = "\n".join(errors_list)

        error = {
            "type": "prompt_outputs_failed_validation",
            "message": "Prompt outputs failed validation",
            "details": errors_list,
            "extra_info": {},
        }

        return (False, error, list(good_outputs), node_errors)

    return (True, None, list(good_outputs), node_errors)


def validate_prompt(prompt):
    from execution import full_type_name, validate_inputs

    result = get_outputs(prompt)

    if result[0] is not True:
        return result

    outputs = result[1]

    good_outputs = set()
    errors = []
    node_errors = {}
    validated = {}
    for o in outputs:
        valid = False
        reasons = []
        try:
            m = validate_inputs(prompt, o, validated)
            valid = m[0]
            reasons = m[1]
        except Exception as ex:
            typ, _, tb = sys.exc_info()
            valid = False
            exception_type = full_type_name(typ)
            reasons = [
                {
                    "type": "exception_during_validation",
                    "message": "Exception when validating node",
                    "details": str(ex),
                    "extra_info": {
                        "exception_type": exception_type,
                        "traceback": traceback.format_tb(tb),
                    },
                }
            ]
            validated[o] = (False, reasons, o)

        if valid is True:
            good_outputs.add(o)
        else:
            logging.error(f"Failed to validate prompt for output {o}:")
            if len(reasons) > 0:
                logging.error("* (prompt):")
                for reason in reasons:
                    logging.error(f"  - {reason['message']}: {reason['details']}")
            errors += [(o, reasons)]
            for node_id, result in validated.items():
                valid = result[0]
                reasons = result[1]
                # If a node upstream has errors, the nodes downstream will also
                # be reported as invalid, but there will be no errors attached.
                # So don't return those nodes as having errors in the response.
                if valid is not True and len(reasons) > 0:
                    if node_id not in node_errors:
                        class_type = prompt[node_id]["class_type"]
                        node_errors[node_id] = {
                            "errors": reasons,
                            "dependent_outputs": [],
                            "class_type": class_type,
                        }
                        logging.error(f"* {class_type} {node_id}:")
                        for reason in reasons:
                            logging.error(
                                f"  - {reason['message']}: {reason['details']}"
                            )
                    node_errors[node_id]["dependent_outputs"].append(o)
            logging.error("Output will be ignored")

    if len(good_outputs) == 0:
        errors_list = []
        for o, errors in errors:
            for error in errors:
                errors_list.append(f"{error['message']}: {error['details']}")
        errors_list = "\n".join(errors_list)

        error = {
            "type": "prompt_outputs_failed_validation",
            "message": "Prompt outputs failed validation",
            "details": errors_list,
            "extra_info": {},
        }

        return (False, error, list(good_outputs), node_errors)

    return (True, None, list(good_outputs), node_errors)


def recursive_execute_og(
    server,
    prompt,
    outputs,
    current_item,
    extra_data,
    executed,
    prompt_id,
    outputs_ui,
    object_storage,
    execute_fn,
    child_nodes_took,
):
    import logging
    import sys
    import traceback

    import comfy.model_management
    import nodes
    from main import execution

    unique_id = current_item
    inputs = prompt[unique_id]["inputs"]
    class_type = prompt[unique_id]["class_type"]
    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
    if unique_id in outputs:
        return (True, None, None)

    for x in inputs:
        input_data = inputs[x]

        if isinstance(input_data, list):
            input_unique_id = input_data[0]
            output_index = input_data[1]  # noqa: F841
            if input_unique_id not in outputs:
                with timed() as timer:
                    result = execute_fn(
                        server,
                        prompt,
                        outputs,
                        input_unique_id,
                        extra_data,
                        executed,
                        prompt_id,
                        outputs_ui,
                        object_storage,
                    )

                child_nodes_took.append(timer.took)

                if result[0] is not True:
                    # Another node failed further upstream
                    return result

    input_data_all = None
    try:
        input_data_all = execution.get_input_data(
            inputs, class_def, unique_id, outputs, prompt, extra_data
        )
        if server.client_id is not None:
            server.last_node_id = unique_id
            server.send_sync(
                "executing",
                {"node": unique_id, "prompt_id": prompt_id},
                server.client_id,
            )

        obj = object_storage.get((unique_id, class_type), None)
        if obj is None:
            obj = class_def()
            object_storage[(unique_id, class_type)] = obj

        output_data, output_ui = execution.get_output_data(obj, input_data_all)
        outputs[unique_id] = output_data
        if len(output_ui) > 0:
            outputs_ui[unique_id] = output_ui
            if server.client_id is not None:
                server.send_sync(
                    "executed",
                    {"node": unique_id, "output": output_ui, "prompt_id": prompt_id},
                    server.client_id,
                )
    except comfy.model_management.InterruptProcessingException as iex:
        logging.info("Processing interrupted")

        # skip formatting inputs/outputs
        error_details = {
            "node_id": unique_id,
        }

        return (False, error_details, iex)
    except Exception as ex:
        typ, _, tb = sys.exc_info()
        exception_type = execution.full_type_name(typ)
        input_data_formatted = {}
        if input_data_all is not None:
            input_data_formatted = {}
            for name, inputs in input_data_all.items():
                input_data_formatted[name] = [execution.format_value(x) for x in inputs]

        output_data_formatted = {}
        for node_id, node_outputs in outputs.items():
            output_data_formatted[node_id] = [
                [execution.format_value(x) for x in l]
                for l in node_outputs  # noqa: E741
            ]

        logging.error("!!! Exception during processing !!!")
        logging.error(traceback.format_exc())

        error_details = {
            "node_id": unique_id,
            "exception_message": str(ex),
            "exception_type": exception_type,
            "traceback": traceback.format_tb(tb),
            "current_inputs": input_data_formatted,
            "current_outputs": output_data_formatted,
        }
        return (False, error_details, ex)

    executed.add(unique_id)

    return (True, None, None)


def recursive_execute(
    server,
    prompt,
    outputs,
    current_item,
    extra_data,
    executed,
    prompt_id,
    outputs_ui,
    object_storage,
):
    with timed() as timer:
        child_nodes_took = []

        result = recursive_execute_og(
            server,
            prompt,
            outputs,
            current_item,
            extra_data,
            executed,
            prompt_id,
            outputs_ui,
            object_storage,
            recursive_execute,
            child_nodes_took,
        )

    if server.client_id is not None:
        server.send_sync(
            "fal-node-timings",
            {
                "node": current_item,
                "prompt_id": prompt_id,
                "took": timer.took - sum(child_nodes_took),
            },
            server.client_id,
        )

    return result


def _gligen_set_position(self, latent_image_shape, position_params, device):
    import torch

    batch, c, h, w = latent_image_shape
    masks = torch.zeros([self.max_objs], device="cpu")
    boxes = []
    positive_embeddings = []
    for p in position_params:
        x1 = (p[4]) / w
        y1 = (p[3]) / h
        x2 = (p[4] + p[2]) / w
        y2 = (p[3] + p[1]) / h
        masks[len(boxes)] = 1.0
        boxes += [torch.tensor((x1, y1, x2, y2)).unsqueeze(0).to("cpu")]
        positive_embeddings += [p[0].to("cpu")]
    append_boxes = []
    append_conds = []
    if len(boxes) < self.max_objs:
        append_boxes = [torch.zeros([self.max_objs - len(boxes), 4], device="cpu")]
        append_conds = [
            torch.zeros([self.max_objs - len(boxes), self.key_dim], device="cpu")
        ]

    box_out = torch.cat(boxes + append_boxes).unsqueeze(0).repeat(batch, 1, 1)
    masks = masks.unsqueeze(0).repeat(batch, 1)
    conds = (
        torch.cat(positive_embeddings + append_conds).unsqueeze(0).repeat(batch, 1, 1)
    )
    return self._set_position(box_out.to(device), masks.to(device), conds.to(device))


async def ainit_extra_nodes(init_custom_nodes: bool = True):
    from nodes import init_builtin_extra_nodes, init_external_custom_nodes, args

    import_failed = await init_builtin_extra_nodes()

    if init_custom_nodes:
        await init_external_custom_nodes()
    else:
        logging.info("Skipping loading of custom nodes")

    if len(import_failed) > 0:
        logging.warning(
            "WARNING: some comfy_extras/ nodes did not import correctly. This may be because they are missing some dependencies.\n"
        )
        for node in import_failed:
            logging.warning("IMPORT FAILED: {}".format(node))
        logging.warning(
            "\nThis issue might be caused by new missing dependencies added the last time you updated ComfyUI."
        )
        if args.windows_standalone_build:
            logging.warning("Please run the update script: update/update_comfyui.bat")
        else:
            logging.warning("Please do a: pip install -r requirements.txt")
        logging.warning("")

    if import_failed:
        logging.warning("Some nodes failed to import: %s (continuing anyway)", import_failed)

    return import_failed


def init_extra_nodes(init_custom_nodes: bool = True):
    from nodes import init_builtin_extra_nodes, init_external_custom_nodes, args

    import_failed = init_builtin_extra_nodes()

    if init_custom_nodes:
        init_external_custom_nodes()
    else:
        logging.info("Skipping loading of custom nodes")

    if len(import_failed) > 0:
        logging.warning(
            "WARNING: some comfy_extras/ nodes did not import correctly. This may be because they are missing some dependencies.\n"
        )
        for node in import_failed:
            logging.warning("IMPORT FAILED: {}".format(node))
        logging.warning(
            "\nThis issue might be caused by new missing dependencies added the last time you updated ComfyUI."
        )
        if args.windows_standalone_build:
            logging.warning("Please run the update script: update/update_comfyui.bat")
        else:
            logging.warning("Please do a: pip install -r requirements.txt")
        logging.warning("")

    if import_failed:
        logging.warning("Some nodes failed to import: %s (continuing anyway)", import_failed)

    return import_failed


async def send_json(self, event, data, sid=None, callback=None):
    from server import send_socket_catch_exception

    message = {"type": event, "data": data}

    if callback:
        callback(json.dumps(message))
    else:
        if sid is None:
            sockets = list(self.sockets.values())
            for ws in sockets:
                await send_socket_catch_exception(ws.send_json, message)
        elif sid in self.sockets:
            await send_socket_catch_exception(self.sockets[sid].send_json, message)
