import bpy

from ...core import Check, Issue, Severity, contexts, is_local

# Node types that make a tree produce something.
_OUTPUT_TYPES = {
    "OUTPUT_MATERIAL", "OUTPUT_WORLD", "OUTPUT_LIGHT", "OUTPUT_AOV", "OUTPUT_LINESTYLE",
    "COMPOSITE", "OUTPUT_FILE", "VIEWER",
}
# Node types never reported as unused.
_IGNORED_TYPES = {"FRAME", "GROUP_INPUT"}


def unused_nodes(tree):
    """Nodes that do not contribute to any output of ``tree``.

    Returns None when the tree has no output at all (nothing to compare with).
    """
    roots = [
        node for node in tree.nodes
        if node.type in _OUTPUT_TYPES
        or (node.type == "GROUP_OUTPUT" and getattr(node, "is_active_output", True))
    ]
    if not roots:
        return None

    upstream = {}
    for link in tree.links:
        upstream.setdefault(link.to_node.name, []).append(link.from_node)

    used = set()
    stack = list(roots)
    while stack:
        node = stack.pop()
        if node.name in used:
            continue
        used.add(node.name)
        stack.extend(upstream.get(node.name, ()))

    return [node for node in tree.nodes if node.name not in used and node.type not in _IGNORED_TYPES]


def _node_trees():
    """(owner datablock, node tree) of every local node tree in use."""
    for material in bpy.data.materials:
        if is_local(material) and material.users and material.node_tree:
            yield material, material.node_tree
    for world in bpy.data.worlds:
        if is_local(world) and world.users and world.node_tree:
            yield world, world.node_tree
    for group in bpy.data.node_groups:
        if is_local(group) and group.users:
            yield group, group


def _tree_of(owner):
    return owner if isinstance(owner, bpy.types.NodeTree) else owner.node_tree


class UnusedNodesCheck(Check):
    id = "unused_nodes"
    label = "Unused nodes"
    description = "Material, world and node group trees contain no node disconnected from the output."
    category = "Shading"
    contexts = contexts.ASSET
    severity = Severity.WARNING
    fix_description = "Delete the unused nodes"

    def run(self, ctx):
        for owner, tree in _node_trees():
            nodes = unused_nodes(tree)
            if nodes:
                names = [node.name for node in nodes]
                shown = ", ".join(names[:5]) + (" ..." if len(names) > 5 else "")
                yield Issue.for_id(
                    f"{owner.name}: {len(names)} unused node(s): {shown}",
                    owner,
                    data={"nodes": names},
                )

    def fix(self, ctx, issues):
        for issue in issues:
            owner = issue.get_id()
            if owner is None:
                continue
            tree = _tree_of(owner)
            for name in issue.data["nodes"]:
                node = tree.nodes.get(name)
                if node is not None:
                    tree.nodes.remove(node)