from PySide6.QtGui import QUndoCommand

from bedit_core.models import BondCausality, BondConnection, Component, ConnectionID, GraphImplementation

CausalityState = dict[ConnectionID, tuple[BondCausality, bool]]


def causality_state(component: Component) -> CausalityState:
    state = {}
    if not isinstance(component.implementation, GraphImplementation):
        return state
    for connection_id, connection in component.implementation.graph.connections.items():
        if isinstance(connection, BondConnection):
            state[connection_id] = (connection.causality, connection.undesired)
    for child in component.implementation.graph.components.values():
        state.update(causality_state(child))
    return state


class ChangeCausalityCommand(QUndoCommand):
    def __init__(self, document: object, component: Component, inferred_component: Component) -> None:
        super().__init__("Infer causality")
        self.document = document
        self.component = component
        self.old_state = causality_state(component)
        self.new_state = causality_state(inferred_component)

    def redo(self) -> None:
        self._apply(self.component, self.new_state)
        self.document.model_changed.emit(self.document.model)

    def undo(self) -> None:
        self._apply(self.component, self.old_state)
        self.document.model_changed.emit(self.document.model)

    @classmethod
    def _apply(cls, component: Component, state: CausalityState) -> None:
        if not isinstance(component.implementation, GraphImplementation):
            return
        for connection_id, connection in component.implementation.graph.connections.items():
            if isinstance(connection, BondConnection) and connection_id in state:
                connection.causality, connection.undesired = state[connection_id]
        for child in component.implementation.graph.components.values():
            cls._apply(child, state)