from __future__ import annotations

from typing import Set

from ..model import Attribute, Diagram, DiagramElement, Entity
from ..model.sheep_wolf_diagram import SHEEP_WOLF_DIAGRAM


# Directory constants
INPUT_DIR = "/design"
OUTPUT_DIR = "/design"


def _visit_children(node: DiagramElement, visited: Set[str]) -> None:
    if node.id in visited:
        return
    visited.add(node.id)
    
    for link in getattr(node, "links", []):
        target = getattr(link, "target", None)
        if target is None:
            continue
        
        link_id = getattr(link, "id", "")
        
        if isinstance(target, Attribute) and link_id.startswith("bel_"):
            if isinstance(node, Entity) and getattr(target, "parent", None) is None:
                target.parent = node
        
        _visit_children(target, visited)


def preprocess_diagram(diagram: Diagram) -> Diagram:
    for node in diagram.nodes.values():
        if isinstance(node, Entity) and node.magnitude_order != -1:
            visited: Set[str] = set()
            _visit_children(node, visited)
    
    return diagram


def preprocess_sheep_wolf_diagram() -> Diagram:
    return preprocess_diagram(SHEEP_WOLF_DIAGRAM)


if __name__ == "__main__":
    preprocess_sheep_wolf_diagram()