from collections import defaultdict
import time


class CauseGraphEngine:
    """
    Builds a causal graph from system events.

    This is NOT execution logic.
    It is a post-hoc + live reasoning layer.
    """

    def __init__(self):
        # node -> causes
        self.graph = defaultdict(list)

        # event timeline
        self.events = []

        # last decision context
        self.last_context = {}

        # bus reference (set by attach_bus)
        self.bus = None

    # -------------------------------------------------
    # BUS BINDING
    # -------------------------------------------------
    def attach_bus(self, bus):
        """Bind to the event bus. Does NOT subscribe — caller does that."""
        self.bus = bus

    def emit_debug(self, data):
        """Emit a debug trace event back to the bus. Use sparingly to avoid
        recursion: subscribers of `debug` events MUST NOT call emit_debug."""
        if self.bus is None:
            return
        self.bus.emit(
            "debug",
            node=None,
            action=None,
            state="trace",
            meta=data,
        )

    # -------------------------------------------------
    # INGEST EVENTS
    # -------------------------------------------------
    def ingest(self, event: dict):
        self.events.append(event)

        etype = event.get("type")
        node = event.get("node")
        action = event.get("action")

        # store context for causal linking
        if etype == "task_update":
            self._handle_task(event)
        elif etype == "cache_update":
            self._handle_cache(event)
        elif etype == "failure":
            self._handle_failure(event)
        elif etype == "node_update":
            self._handle_node(event)

    # -------------------------------------------------
    # TASK EXECUTION CAUSALITY
    # -------------------------------------------------
    def _handle_task(self, event):
        node = event.get("node")
        action = event.get("action")
        state = event.get("state")

        context = {
            "time": time.time(),
            "node": node,
            "action": action,
            "state": state,
            "reason": "scheduler_assignment",
        }

        self.last_context[(node, action)] = context

        self.graph[(node, action)].append(context)

    # -------------------------------------------------
    # CACHE CAUSALITY
    # -------------------------------------------------
    def _handle_cache(self, event):
        node = event.get("node")
        action = event.get("action")

        context = {
            "time": time.time(),
            "node": node,
            "action": action,
            "reason": "cache_hit_or_miss",
        }

        self.graph[(node, action)].append(context)

    # -------------------------------------------------
    # FAILURE CAUSALITY
    # -------------------------------------------------
    def _handle_failure(self, event):
        node = event.get("node")
        action = event.get("action")

        prev = self.last_context.get((node, action), {})

        context = {
            "time": time.time(),
            "node": node,
            "action": action,
            "reason": "execution_failure",
            "parent_context": prev,
        }

        self.graph[(node, action)].append(context)

    # -------------------------------------------------
    # NODE STATE CAUSALITY
    # -------------------------------------------------
    def _handle_node(self, event):
        node = event.get("node")

        context = {
            "time": time.time(),
            "node": node,
            "reason": "node_state_update",
            "state": event.get("node_state", {}),
        }

        self.graph[(node, None)].append(context)

    # -------------------------------------------------
    # QUERY INTERFACE
    # -------------------------------------------------
    def explain(self, node, action=None):
        """
        Returns causal chain for a given execution unit.

        If action is None, returns ALL events for the given node
        (any action). Otherwise returns events for the specific
        (node, action) pair.
        """
        if action is not None:
            return self.graph.get((node, action), [])

        # Node-only query: scan all keys starting with this node
        results = []
        for (n, a), events in self.graph.items():
            if n == node:
                results.extend(events)
        # Sort by time
        results.sort(key=lambda e: e.get("time", 0))
        return results

    def full_trace(self):
        return dict(self.graph)