"""
Subworkflow: loads an inner workflow (UI or API format), infers its I/O from
Subworkflow Input / Subworkflow Output nodes, and executes it as a subgraph expansion.
"""
import logging
import weakref

from comfy_api.latest import io
from .debug_utils import configure_logger

from .workflow_utils import (
    list_workflow_files,
    load_workflow_file,
    load_workflow_url,
    build_expansion,
    build_modifier_source_expansion,
    apply_control_after_generate,
    validate_workflow_nodes_installed,
    get_workflow_interface,
)

log = configure_logger("ComfyUI-Subworkflow")


class _UnboundedOutputTypes:
    """Returns '*' for any slot index so validation never IndexErrors on dynamic outputs."""
    def __getitem__(self, idx):
        return "*"
    def __len__(self):
        return 0
    def __iter__(self):
        return iter(())


class _UnboundedFalseSeq:
    """Infinite False iterator so merge_result_data zips correctly for any output count."""
    def __getitem__(self, idx):
        return False
    def __len__(self):
        return 0
    def __iter__(self):
        while True:
            yield False


def _apply_primitive_overrides(data: dict, kwargs: dict) -> dict:
    """
    For INT/FLOAT inputs where the outer slot is not connected and the override
    widget has its switch enabled, substitute the widget value so build_expansion
    uses it instead of the inner fallback node.

    The frontend sends a single widget named swf_override_i whose value is a dict
    {"use": bool, "val": number}.
    """
    info = get_workflow_interface(data)
    effective = dict(kwargs)
    for i, inp_info in enumerate(info["inputs"]):
        if inp_info.get("type") not in ("INT", "FLOAT"):
            continue
        key = f"swf_in_{i}"
        if effective.get(key) is not None:
            continue  # connected node takes priority
        override = effective.get(f"swf_override_{i}")
        if isinstance(override, dict) and override.get("use"):
            effective[key] = override.get("val")
    return effective


class BaseSubworkflow(io.ComfyNode):
    """
    Shared execution behavior for subworkflow nodes that load workflow JSON from
    different sources.
    """

    @classmethod
    def validate_inputs(cls, **kwargs):
        # Accept any incoming types for dynamic swf_in_* inputs.
        return True

    @classmethod
    def check_lazy_status(cls, **kwargs):
        missing_inputs = [
            name for name, value in kwargs.items()
            if name.startswith("swf_in_") and value is None
        ]
        if missing_inputs:
            log.debug("[Subworkflow] waiting for dynamic input(s) %s", missing_inputs)
        return missing_inputs

    @classmethod
    def fingerprint_inputs(cls, *args, **kwargs):
        # The inner workflow can change between runs by file reload or by cached
        # control-after-generate mutations, even when outer inputs are unchanged.
        return float("NaN")

    _loaded_workflows: dict[str, dict] = {}
    _loaded_workflow_instances: dict[str, tuple[str, dict]] = {}
    _prompt_workflows = weakref.WeakKeyDictionary()

    @classmethod
    def _source_label(cls, **kwargs) -> str:
        raise NotImplementedError

    @classmethod
    def _source_cache_key(cls, **kwargs) -> str:
        raise NotImplementedError

    @classmethod
    def _load_source(cls, **kwargs) -> dict:
        raise NotImplementedError

    @classmethod
    def _get_workflow_data(
        cls,
        reload_each_execution: bool,
        *,
        instance_id: str | None = None,
        dynprompt=None,
        **kwargs,
    ) -> dict:
        cache_key = cls._source_cache_key(**kwargs)
        label = cls._source_label(**kwargs)

        prompt_cache = None
        prompt_cache_key = None
        if dynprompt is not None and instance_id is not None:
            try:
                prompt_cache = cls._prompt_workflows.setdefault(dynprompt, {})
                prompt_cache_key = (str(instance_id), cache_key)
                if prompt_cache_key in prompt_cache:
                    log.debug(
                        "[Subworkflow] reusing prompt-scoped inner workflow %r for node %s",
                        label,
                        instance_id,
                    )
                    return prompt_cache[prompt_cache_key]
            except TypeError:
                # Some test and mock prompt objects may not support weak references.
                prompt_cache = None
                prompt_cache_key = None

        instance_key = str(instance_id) if instance_id is not None else None
        instance_cached = (
            cls._loaded_workflow_instances.get(instance_key)
            if instance_key is not None
            else None
        )
        instance_cache_hit = (
            instance_cached is not None and instance_cached[0] == cache_key
        )

        if reload_each_execution or (
            instance_key is not None and not instance_cache_hit
        ) or (
            instance_key is None and cache_key not in cls._loaded_workflows
        ):
            reason = "reload_each_execution enabled" if reload_each_execution else "cache miss"
            log.info(
                "[Subworkflow] loading inner workflow %r for node %s (%s)",
                label,
                instance_id or "shared source",
                reason,
            )
            data = cls._load_source(**kwargs)
            if instance_key is not None:
                if not reload_each_execution:
                    cls._loaded_workflow_instances[instance_key] = (cache_key, data)
            else:
                cls._loaded_workflows[cache_key] = data
        elif instance_key is not None:
            data = instance_cached[1]
        else:
            log.debug("[Subworkflow] reusing cached inner workflow %r", label)
            data = cls._loaded_workflows[cache_key]

        if prompt_cache is not None and prompt_cache_key is not None:
            prompt_cache[prompt_cache_key] = data
        return data

    @classmethod
    def _execute_source(cls, reload_each_execution=True, **kwargs):
        label = cls._source_label(**kwargs)
        log.debug(
            "[Subworkflow] executing inner workflow %r (reload_each_execution=%s, dynamic_inputs=%s)",
            label,
            reload_each_execution,
            sorted(k for k in kwargs if k.startswith("swf_in_")),
        )
        data = cls._get_workflow_data(
            reload_each_execution,
            instance_id=cls.hidden.unique_id,
            dynprompt=cls.hidden.dynprompt,
            **kwargs,
        )
        validate_workflow_nodes_installed(data)

        effective_inputs = _apply_primitive_overrides(data, kwargs)
        output_refs, graph = build_expansion(data, effective_inputs)
        return cls._finalize_execution(output_refs, graph, data, reload_each_execution)

    @classmethod
    def _finalize_execution(cls, output_refs, graph, data: dict, reload_each_execution: bool):
        if not reload_each_execution:
            apply_control_after_generate(data)
        return io.NodeOutput(*output_refs, expand=graph.finalize())


def _subworkflow_outputs():
    return []


def _reload_input():
    return io.Boolean.Input(
        "reload_each_execution",
        display_name="at execution",
        default=True,
        label_on="reload",
        label_off="keep loaded",
    )


class Subworkflow(BaseSubworkflow):
    """
    Selects a saved workflow, exposes its Subworkflow Input nodes as inputs and
    its Subworkflow Output nodes as outputs, then executes the inner workflow as
    a transparent subgraph when the outer workflow runs.
    """

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SWF_Subworkflow",
            display_name="Subworkflow",
            category="Subworkflow",
            description=(
                "Executes a selected workflow as an expandable subworkflow, "
                "using Subworkflow Input and Subworkflow Output boundary nodes."
            ),
            inputs=[
                io.Combo.Input("workflow", options=list_workflow_files()),
                _reload_input(),
            ],
            hidden=[io.Hidden.unique_id, io.Hidden.dynprompt],
            outputs=_subworkflow_outputs(),
            enable_expand=True,
            accept_all_inputs=True,
        )

    @classmethod
    def _source_label(cls, workflow: str, **kwargs) -> str:
        return workflow

    @classmethod
    def _source_cache_key(cls, workflow: str, **kwargs) -> str:
        return f"file:{workflow}"

    @classmethod
    def _load_source(cls, workflow: str, **kwargs) -> dict:
        if not workflow:
            raise ValueError("No workflow selected. Choose a workflow file from the dropdown.")
        try:
            return load_workflow_file(workflow)
        except FileNotFoundError:
            raise ValueError(f"Workflow file not found: {workflow!r}") from None

    @classmethod
    def execute(cls, workflow: str, reload_each_execution=True, **kwargs):
        return cls._execute_source(
            workflow=workflow,
            reload_each_execution=reload_each_execution,
            **kwargs,
        )


class SubworkflowFromURL(BaseSubworkflow):
    """
    Loads a workflow JSON document from an HTTP(S) URL and executes it as a
    transparent subgraph.
    """

    @classmethod
    def define_schema(cls):
        log.debug("[Subworkflow] SubworkflowFromURL define_schema called")
        return io.Schema(
            node_id="SWF_SubworkflowFromURL",
            display_name="Subworkflow (from URL)",
            category="Subworkflow",
            description=(
                "Executes a workflow loaded from an HTTP(S) URL as an expandable "
                "subworkflow, using Subworkflow Input and Subworkflow Output "
                "boundary nodes."
            ),
            inputs=[
                io.String.Input("url", multiline=False, default=""),
                io.Boolean.Input(
                    "verify_ssl",
                    display_name="verify SSL",
                    default=True,
                    label_on="verify",
                    label_off="skip",
                ),
                _reload_input(),
            ],
            hidden=[io.Hidden.unique_id, io.Hidden.dynprompt],
            outputs=_subworkflow_outputs(),
            enable_expand=True,
            accept_all_inputs=True,
        )

    @classmethod
    def _source_label(cls, url: str, **kwargs) -> str:
        log.debug("[Subworkflow] SubworkflowFromURL _source_label url=%r", url)
        return url

    @classmethod
    def _source_cache_key(cls, url: str, verify_ssl: bool = True, **kwargs) -> str:
        log.debug(
            "[Subworkflow] SubworkflowFromURL _source_cache_key url=%r stripped=%r verify_ssl=%s",
            url,
            url.strip(),
            verify_ssl,
        )
        return f"url:{url.strip()}:verify_ssl:{bool(verify_ssl)}"

    @classmethod
    def _load_source(cls, url: str, verify_ssl: bool = True, **kwargs) -> dict:
        log.debug(
            "[Subworkflow] SubworkflowFromURL _load_source url=%r verify_ssl=%s kwargs_keys=%s",
            url,
            verify_ssl,
            sorted(kwargs.keys()),
        )
        if not url:
            log.warning("[Subworkflow] SubworkflowFromURL _load_source missing URL")
            raise ValueError("No workflow URL specified.")
        return load_workflow_url(url.strip(), verify_ssl=bool(verify_ssl))

    @classmethod
    def execute(cls, url: str, verify_ssl=True, reload_each_execution=True, **kwargs):
        log.debug(
            "[Subworkflow] SubworkflowFromURL execute called url=%r verify_ssl=%s reload_each_execution=%s dynamic_inputs=%s",
            url,
            verify_ssl,
            reload_each_execution,
            sorted(k for k in kwargs if k.startswith("swf_in_")),
        )
        return cls._execute_source(
            url=url,
            verify_ssl=verify_ssl,
            reload_each_execution=reload_each_execution,
            **kwargs,
        )


class SubworkflowModifierSource(BaseSubworkflow):
    """
    Exposes a modifier-source slot from an inner workflow as a separate outer node.
    This breaks the outer dependency cycle by moving the pre-modifier value onto
    its own node, while the main Subworkflow keeps only the modifier input.
    """

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SWF_SubworkflowModifierSource",
            display_name="Subworkflow Modifier Source",
            category="Subworkflow",
            description=(
                "Exposes modifier source slots from another Subworkflow node in "
                "this workflow. Use this together with that Subworkflow node to "
                "break outer modifier dependency loops."
            ),
            inputs=[
                io.Combo.Input(
                    "target_subworkflow",
                    display_name="Subworkflow node",
                    options=[""],
                ),
            ],
            hidden=[io.Hidden.unique_id, io.Hidden.dynprompt],
            outputs=_subworkflow_outputs(),
            enable_expand=True,
            accept_all_inputs=True,
        )

    @classmethod
    def _target_node_id(cls, target_subworkflow: str) -> str:
        return str(target_subworkflow or "").split(" | ", 1)[0].strip()

    @classmethod
    def _resolve_target(cls, target_subworkflow: str):
        target_id = cls._target_node_id(target_subworkflow)
        if not target_id:
            raise ValueError("No Subworkflow node selected.")

        dynprompt = cls.hidden.dynprompt
        if dynprompt is None or not dynprompt.has_node(target_id):
            raise ValueError(
                f"Selected Subworkflow node {target_id!r} is not present in the queued workflow."
            )

        target = dynprompt.get_node(target_id)
        inputs = target.get("inputs") or {}
        class_type = target.get("class_type")
        if class_type == "SWF_Subworkflow":
            workflow = inputs.get("workflow")
            if not isinstance(workflow, str) or not workflow:
                raise ValueError(
                    f"Selected Subworkflow node {target_id!r} has no workflow selected."
                )
            source_class = Subworkflow
            source_kwargs = {"workflow": workflow}
        elif class_type == "SWF_SubworkflowFromURL":
            url = inputs.get("url")
            verify_ssl = inputs.get("verify_ssl", True)
            if not isinstance(url, str) or not url:
                raise ValueError(
                    f"Selected Subworkflow node {target_id!r} has no workflow URL."
                )
            if not isinstance(verify_ssl, bool):
                raise ValueError(
                    f"Selected Subworkflow node {target_id!r} has an invalid verify-SSL setting."
                )
            source_class = SubworkflowFromURL
            source_kwargs = {"url": url, "verify_ssl": verify_ssl}
        else:
            raise ValueError(
                f"Selected node {target_id!r} is not a Subworkflow node."
            )

        reload_each_execution = inputs.get("reload_each_execution", True)
        if not isinstance(reload_each_execution, bool):
            raise ValueError(
                f"Selected Subworkflow node {target_id!r} has an invalid at-execution setting."
            )
        return target_id, source_class, source_kwargs, reload_each_execution

    @classmethod
    def execute(cls, target_subworkflow: str, **kwargs):
        target_id, source_class, source_kwargs, reload_each_execution = (
            cls._resolve_target(target_subworkflow)
        )
        data = source_class._get_workflow_data(
            reload_each_execution,
            instance_id=target_id,
            dynprompt=cls.hidden.dynprompt,
            **source_kwargs,
        )
        validate_workflow_nodes_installed(data)
        output_refs, graph = build_modifier_source_expansion(data, kwargs)
        # The selected Subworkflow owns control-after-generate mutations.
        return io.NodeOutput(*output_refs, expand=graph.finalize())


class SubworkflowModifierSourceFromURL(BaseSubworkflow):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="SWF_SubworkflowModifierSourceFromURL",
            display_name="Subworkflow Modifier Source (from URL)",
            category="Subworkflow",
            is_deprecated=True,
            description=(
                "Deprecated compatibility node. Use Subworkflow Modifier Source "
                "and select a Subworkflow (from URL) node instead."
            ),
            inputs=[
                io.String.Input("url", multiline=False, default=""),
                io.Boolean.Input(
                    "verify_ssl",
                    display_name="verify SSL",
                    default=True,
                    label_on="verify",
                    label_off="skip",
                ),
            ],
            outputs=_subworkflow_outputs(),
            enable_expand=True,
            accept_all_inputs=True,
        )

    @classmethod
    def _source_label(cls, url: str, **kwargs) -> str:
        return url

    @classmethod
    def _source_cache_key(cls, url: str, verify_ssl: bool = True, **kwargs) -> str:
        return f"url:{url.strip()}:verify_ssl:{bool(verify_ssl)}"

    @classmethod
    def _load_source(cls, url: str, verify_ssl: bool = True, **kwargs) -> dict:
        if not url:
            raise ValueError("No workflow URL specified.")
        return load_workflow_url(url.strip(), verify_ssl=bool(verify_ssl))

    @classmethod
    def execute(cls, url: str, verify_ssl: bool, **kwargs):
        reload_each_execution = True
        data = cls._get_workflow_data(
            reload_each_execution,
            url=url,
            verify_ssl=verify_ssl,
            **kwargs,
        )
        validate_workflow_nodes_installed(data)
        output_refs, graph = build_modifier_source_expansion(data, kwargs)
        return cls._finalize_execution(output_refs, graph, data, reload_each_execution)


_DYNAMIC_RETURN = _UnboundedOutputTypes()
for _cls in (Subworkflow, SubworkflowFromURL, SubworkflowModifierSource, SubworkflowModifierSourceFromURL):
    # Pre-set so GET_SCHEMA() doesn't overwrite with the empty list that outputs=[] produces.
    # Returning '*' for any slot index lets validation accept connections to any output slot.
    _cls._RETURN_TYPES = _DYNAMIC_RETURN
    _cls._RETURN_NAMES = []
    _cls._OUTPUT_IS_LIST = _UnboundedFalseSeq()
    _cls._OUTPUT_TOOLTIPS = []