from .edges import Edge, START
from .nodes import Node
from .states import State, Shared
from typing import Type, TypeVar, Generic, Callable, Sequence, Coroutine, Tuple
from collections import defaultdict
from deepdiff import DeepDiff
import asyncio

T = TypeVar('T', bound=State)
S = TypeVar('S', bound=Shared)

class GraphExecutor(Generic[T, S]):

    edges: list[Edge[T, S]]

    def __init__(self, edges: list[Edge[T, S]]):
        self.edges = edges
        self._index: dict[Node[T, S] | Type[START], list[Edge[T, S]]] = defaultdict(list[Edge[T, S]])

        for edge in self.edges:
            if isinstance(edge.source, list):
                for source in edge.source:
                    self._index[source].append(edge)
            else:
                self._index[edge.source].append(edge)


    async def __call__(self, initial_state: T, initial_shared: S) -> Tuple[T, S]:
        state: T = initial_state
        shared: S = initial_shared
        current_nodes: Sequence[Node[T, S] | Type[START]] = [START]

        while True:

            edges: list[Edge[T, S]] = []

            for current_node in current_nodes:

                # Find the edge corresponding to the current node
                edges.extend(self._index[current_node])


            next_nodes: list[Node[T, S]] = [
                n for edge in edges 
                if isinstance((n := edge.next(state, shared)), Node)
            ] # Only one execution of next filtering "END"s


            if not next_nodes:
                break # END

            
            parallel_tasks: list[Callable[[T, S], Coroutine[None, None, T]]] = []
            # Determine the next node using the edge's next function
            for next_node in next_nodes:
                
                parallel_tasks.append(next_node.run)

            # Run parallel

            tasks: list[asyncio.Task[T]] = []

            async with asyncio.TaskGroup() as tg:
                for task in parallel_tasks:
                    
                    state_copy: T = state.model_copy(deep=True)

                    tasks.append(tg.create_task(task(state_copy, shared)))

            result_states: list[T] = [task.result() for task in tasks]

            state = self.merge_states(state, result_states)

            current_nodes = next_nodes


        return state, shared

        


    def merge_states(self, current_state: T, result_states: list[T]) -> T:
            
        result_dicts = [state.model_dump() for state in result_states]
        current_dict = current_state.model_dump()

        diffs = [DeepDiff(current_dict, result_dict) for result_dict in result_dicts]

        print(diffs)

        state: T = T.model_validate(current_dict) # type: ignore

        return state