import importlib.util
import logging
import pathlib
import sys
import types
import unittest


REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]


def _load_workflow_utils_module():
    package_name = "swf_utils_isolated_test"
    package = types.ModuleType(package_name)
    package.__path__ = [str(REPO_ROOT)]
    sys.modules[package_name] = package

    graph_utils = types.ModuleType("comfy_execution.graph_utils")
    graph_utils.GraphBuilder = object
    comfy_execution = types.ModuleType("comfy_execution")
    comfy_execution.graph_utils = graph_utils
    sys.modules["comfy_execution"] = comfy_execution
    sys.modules["comfy_execution.graph_utils"] = graph_utils

    debug_utils = types.ModuleType(f"{package_name}.debug_utils")
    debug_utils.configure_logger = lambda name: logging.getLogger(name)
    sys.modules[debug_utils.__name__] = debug_utils

    module_name = f"{package_name}.workflow_utils"
    spec = importlib.util.spec_from_file_location(
        module_name,
        REPO_ROOT / "workflow_utils.py",
    )
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


workflow_utils = _load_workflow_utils_module()


class FakeNode:
    def __init__(self, **inputs):
        self.inputs = inputs


class FakeGraph:
    def __init__(self, nodes):
        self.nodes = nodes


class ModifierSourcePruningTests(unittest.TestCase):
    def test_keeps_only_selected_output_ancestors(self):
        graph = FakeGraph({
            "source": FakeNode(value=1),
            "modifier_value": FakeNode(value=["source", 0]),
            "sampler": FakeNode(model=["modifier_value", 0]),
            "save": FakeNode(images=["sampler", 0]),
            "selected": FakeNode(nested={"value": ["modifier_value", 0]}),
        })

        workflow_utils._prune_graph_to_output_refs(
            graph,
            [["selected", 0]],
        )

        self.assertEqual(
            set(graph.nodes),
            {"source", "modifier_value", "selected"},
        )


if __name__ == "__main__":
    unittest.main()