"""DAG editor session — manages graph state with undo/redo."""

from __future__ import annotations

from dataclasses import dataclass, field

from orxaq.kernel.types import CausalEdge, CausalRole, CycleError, Entity, VariableType
from orxaq.kernel.world_dag import WorldDAG


@dataclass
class EditAction:
    """A recorded graph edit for undo/redo."""

    action_type: str  # "add_node", "remove_node", "add_edge", "remove_edge"
    data: dict = field(default_factory=dict)


class DAGEditorSession:
    """Interactive editing session with undo/redo stack.

    Wraps a WorldDAG with state tracking, enabling the editor
    UI to reverse any change.
    """

    def __init__(self, dag: WorldDAG | None = None) -> None:
        self.dag = dag or WorldDAG()
        self._undo_stack: list[EditAction] = []
        self._redo_stack: list[EditAction] = []
        self._max_history = 100

    def add_node(
        self,
        entity_id: str,
        name: str = "",
        description: str = "",
        causal_role: str = "exogenous",
        variable_type: str = "static",
        domain: str = "",
    ) -> dict:
        """Add a node, returning status dict."""
        if self.dag.get_entity(entity_id) is not None:
            return {"error": f"node '{entity_id}' already exists", "status": "conflict"}

        try:
            role = CausalRole(causal_role)
        except ValueError:
            role = CausalRole.EXOGENOUS

        try:
            vtype = VariableType(variable_type)
        except ValueError:
            vtype = VariableType.STATIC

        entity = Entity(
            entity_id=entity_id,
            name=name or entity_id,
            description=description,
            causal_role=role,
            variable_type=vtype,
            domain=domain,
        )
        self.dag.add_entity(entity)

        self._push_undo(
            EditAction(
                action_type="add_node",
                data={"entity_id": entity_id, "entity": _entity_to_dict(entity)},
            )
        )
        return {"status": "created", "node": entity_id}

    def remove_node(self, entity_id: str) -> dict:
        """Remove a node and its edges, returning status dict."""
        entity = self.dag.get_entity(entity_id)
        if entity is None:
            return {"error": f"node '{entity_id}' not found", "status": "not_found"}

        connected_edges = []
        for edge in self.dag._edges.values():
            if edge.source == entity_id or edge.target == entity_id:
                connected_edges.append(_edge_to_dict(edge))

        self.dag.remove_entity(entity_id)

        self._push_undo(
            EditAction(
                action_type="remove_node",
                data={
                    "entity": _entity_to_dict(entity),
                    "connected_edges": connected_edges,
                },
            )
        )
        return {"status": "deleted", "node": entity_id}

    def add_edge(
        self,
        source: str,
        target: str,
        mechanism: str = "",
        confidence: float = 0.5,
    ) -> dict:
        """Add an edge, returning status dict."""
        edge = CausalEdge(
            source=source,
            target=target,
            mechanism=mechanism,
            confidence=confidence,
        )
        try:
            self.dag.add_edge(edge)
        except CycleError as e:
            return {"error": str(e), "status": "cycle"}
        except ValueError as e:
            return {"error": str(e), "status": "invalid"}

        self._push_undo(
            EditAction(
                action_type="add_edge",
                data={"edge": _edge_to_dict(edge)},
            )
        )
        return {"status": "created", "edge": f"{source}->{target}"}

    def remove_edge(self, source: str, target: str) -> dict:
        """Remove an edge, returning status dict."""
        edge_id = f"{source}\u2192{target}"
        edge = self.dag._edges.get(edge_id)
        if edge is None:
            return {"error": f"edge '{source}->{target}' not found", "status": "not_found"}

        self.dag._children_map.get(source, set()).discard(target)
        self.dag._parents_map.get(target, set()).discard(source)
        del self.dag._edges[edge_id]

        self._push_undo(
            EditAction(
                action_type="remove_edge",
                data={"edge": _edge_to_dict(edge)},
            )
        )
        return {"status": "deleted", "edge": f"{source}->{target}"}

    def undo(self) -> dict:
        """Undo the last action."""
        if not self._undo_stack:
            return {"error": "nothing to undo", "status": "empty"}

        action = self._undo_stack.pop()
        self._redo_stack.append(action)

        if action.action_type == "add_node":
            eid = action.data["entity_id"]
            self.dag.remove_entity(eid)
            return {"status": "undone", "action": "add_node", "node": eid}

        if action.action_type == "remove_node":
            entity_data = action.data["entity"]
            entity = _dict_to_entity(entity_data)
            self.dag.add_entity(entity)
            for ed in action.data.get("connected_edges", []):
                try:
                    self.dag.add_edge(_dict_to_edge(ed))
                except (CycleError, ValueError):
                    pass
            return {"status": "undone", "action": "remove_node"}

        if action.action_type == "add_edge":
            ed = action.data["edge"]
            eid = f"{ed['source']}\u2192{ed['target']}"
            if eid in self.dag._edges:
                self.dag._children_map.get(ed["source"], set()).discard(ed["target"])
                self.dag._parents_map.get(ed["target"], set()).discard(ed["source"])
                del self.dag._edges[eid]
            return {"status": "undone", "action": "add_edge"}

        if action.action_type == "remove_edge":
            ed = action.data["edge"]
            try:
                self.dag.add_edge(_dict_to_edge(ed))
            except (CycleError, ValueError):
                pass
            return {"status": "undone", "action": "remove_edge"}

        return {"error": f"unknown action: {action.action_type}"}

    def redo(self) -> dict:
        """Redo the last undone action."""
        if not self._redo_stack:
            return {"error": "nothing to redo", "status": "empty"}

        action = self._redo_stack.pop()
        self._undo_stack.append(action)

        if action.action_type == "add_node":
            ed = action.data["entity"]
            entity = _dict_to_entity(ed)
            self.dag.add_entity(entity)
            return {"status": "redone", "action": "add_node"}

        if action.action_type == "remove_node":
            eid = action.data["entity"]["entity_id"]
            self.dag.remove_entity(eid)
            return {"status": "redone", "action": "remove_node"}

        if action.action_type == "add_edge":
            ed = action.data["edge"]
            try:
                self.dag.add_edge(_dict_to_edge(ed))
            except (CycleError, ValueError):
                pass
            return {"status": "redone", "action": "add_edge"}

        if action.action_type == "remove_edge":
            ed = action.data["edge"]
            eid = f"{ed['source']}\u2192{ed['target']}"
            if eid in self.dag._edges:
                self.dag._children_map.get(ed["source"], set()).discard(ed["target"])
                self.dag._parents_map.get(ed["target"], set()).discard(ed["source"])
                del self.dag._edges[eid]
            return {"status": "redone", "action": "remove_edge"}

        return {"error": f"unknown action: {action.action_type}"}

    @property
    def can_undo(self) -> bool:
        return len(self._undo_stack) > 0

    @property
    def can_redo(self) -> bool:
        return len(self._redo_stack) > 0

    @property
    def history_length(self) -> int:
        return len(self._undo_stack)

    def _push_undo(self, action: EditAction) -> None:
        self._undo_stack.append(action)
        self._redo_stack.clear()
        if len(self._undo_stack) > self._max_history:
            self._undo_stack.pop(0)

    def snapshot(self) -> dict:
        """Current graph state as a serializable dict."""
        nodes = []
        for eid, entity in sorted(self.dag._entities.items()):
            nodes.append(_entity_to_dict(entity))
        edges = []
        for edge in self.dag._edges.values():
            edges.append(_edge_to_dict(edge))
        return {
            "nodes": nodes,
            "edges": edges,
            "node_count": len(nodes),
            "edge_count": len(edges),
            "can_undo": self.can_undo,
            "can_redo": self.can_redo,
        }

    def load_snapshot(self, data: dict) -> None:
        """Load graph from a snapshot dict."""
        self.dag = WorldDAG()
        for nd in data.get("nodes", []):
            self.dag.add_entity(_dict_to_entity(nd))
        for ed in data.get("edges", []):
            try:
                self.dag.add_edge(_dict_to_edge(ed))
            except (CycleError, ValueError):
                pass
        self._undo_stack.clear()
        self._redo_stack.clear()


def _entity_to_dict(entity: Entity) -> dict:
    return {
        "entity_id": entity.entity_id,
        "name": entity.name,
        "description": entity.description,
        "causal_role": entity.causal_role.value,
        "variable_type": entity.variable_type.value,
        "domain": entity.domain,
    }


def _dict_to_entity(d: dict) -> Entity:
    try:
        role = CausalRole(d.get("causal_role", "exogenous"))
    except ValueError:
        role = CausalRole.EXOGENOUS
    try:
        vtype = VariableType(d.get("variable_type", "static"))
    except ValueError:
        vtype = VariableType.STATIC

    return Entity(
        entity_id=d["entity_id"],
        name=d.get("name", d["entity_id"]),
        description=d.get("description", ""),
        causal_role=role,
        variable_type=vtype,
        domain=d.get("domain", ""),
    )


def _edge_to_dict(edge: CausalEdge) -> dict:
    return {
        "source": edge.source,
        "target": edge.target,
        "mechanism": edge.mechanism,
        "confidence": edge.confidence,
        "soliton_score": edge.soliton_score,
    }


def _dict_to_edge(d: dict) -> CausalEdge:
    return CausalEdge(
        source=d["source"],
        target=d["target"],
        mechanism=d.get("mechanism", ""),
        confidence=d.get("confidence", 0.5),
        soliton_score=d.get("soliton_score", 0.5),
    )
