from __future__ import annotations

from pydantic import BaseModel, Field

from sir import EdgeKind, NodeKind, SemanticNode, SirSnapshot


class FormSemantics(BaseModel):
    form: SemanticNode
    commands: list[SemanticNode] = Field(default_factory=list)
    elements: list[SemanticNode] = Field(default_factory=list)
    command_handlers: dict[str, SemanticNode] = Field(default_factory=dict)


def form_semantics(snapshot: SirSnapshot) -> list[FormSemantics]:
    nodes = {node.lineage_id: node for node in snapshot.nodes}
    forms = {
        node.lineage_id: FormSemantics(form=node)
        for node in snapshot.nodes
        if node.kind == NodeKind.FORM
    }
    element_children: dict[str, list[SemanticNode]] = {}
    for edge in snapshot.edges:
        form = forms.get(edge.source_lineage)
        target = nodes.get(edge.target_lineage)
        if target is None:
            continue
        if edge.kind == EdgeKind.HAS_ELEMENT and target.kind == NodeKind.FORM_ELEMENT:
            element_children.setdefault(edge.source_lineage, []).append(target)
        if form is None:
            continue
        if edge.kind == EdgeKind.HAS_COMMAND:
            form.commands.append(target)
    for form in forms.values():
        form.elements.extend(_flatten_form_elements(form.form.lineage_id, element_children))
    command_to_form = {
        command.lineage_id: form
        for form in forms.values()
        for command in form.commands
    }
    for edge in snapshot.edges:
        if edge.kind != EdgeKind.HANDLES:
            continue
        form = command_to_form.get(edge.source_lineage)
        handler = nodes.get(edge.target_lineage)
        if form is not None and handler is not None:
            form.command_handlers[edge.source_lineage] = handler
    return sorted(forms.values(), key=lambda item: item.form.qualified_name)


def _flatten_form_elements(root_lineage: str, element_children: dict[str, list[SemanticNode]]) -> list[SemanticNode]:
    result: list[SemanticNode] = []
    seen: set[str] = set()

    def visit(parent_lineage: str) -> None:
        for element in element_children.get(parent_lineage, []):
            if element.lineage_id in seen:
                continue
            seen.add(element.lineage_id)
            result.append(element)
            visit(element.lineage_id)

    visit(root_lineage)
    return result


__all__ = ["FormSemantics", "form_semantics"]