import logging
from collections import defaultdict, deque
from typing import Any, Dict, List

logger = logging.getLogger(__name__)

class DAGExecutor:
    def __init__(self):
        self.nodes: Dict[str, Dict[str, Any]] = {}
        self.edges: Dict[str, List[str]] = defaultdict(list)
        self.in_degree: Dict[str, int] = defaultdict(int)

    def add_node(self, id: str, agent: Any, task: Any) -> None:
        if id in self.nodes:
            raise ValueError(f"Node {id} already exists.")
        self.nodes[id] = {'agent': agent, 'task': task}
        if id not in self.in_degree:
            self.in_degree[id] = 0

    def add_edge(self, from_id: str, to_id: str) -> None:
        if from_id not in self.nodes or to_id not in self.nodes:
            raise ValueError("Both from_id and to_id must exist as nodes before adding an edge.")
        self.edges[from_id].append(to_id)
        self.in_degree[to_id] += 1

    def execute(self) -> Dict[str, Any]:
        queue = deque([node_id for node_id, degree in self.in_degree.items() if degree == 0])
        results: Dict[str, Any] = {}
        
        if not queue and self.nodes:
            raise ValueError("Graph contains a cycle or has no starting nodes.")
            
        executed_nodes = 0
        
        while queue:
            current_id = queue.popleft()
            node = self.nodes[current_id]
            agent = node['agent']
            task = node['task']
            
            # Pass all available results to the agent (this includes its predecessors)
            try:
                if hasattr(agent, 'execute'):
                    result = agent.execute(task, results)
                elif callable(agent):
                    result = agent(task, results)
                else:
                    raise TypeError(f"Agent for node {current_id} is not callable and has no execute method.")
                
                results[current_id] = result
                executed_nodes += 1
                
                for neighbor in self.edges[current_id]:
                    self.in_degree[neighbor] -= 1
                    if self.in_degree[neighbor] == 0:
                        queue.append(neighbor)
            except Exception as e:
                # We want to catch the error but re-raise to reflect failure
                logger.error(f"Execution failed at node {current_id}: {e}")
                raise RuntimeError(f"Execution failed at node {current_id}") from e
                
        if executed_nodes != len(self.nodes):
            raise ValueError("Graph contains a cycle; not all nodes could be executed.")
            
        return results