from collections import defaultdict

from federation.graph import FederatedGraph
from hiku.engine import (
    InitOptions,
    Query,
    Context,
)
from hiku.executors.queue import Queue


class Engine:
    def __init__(self, executor):
        self.executor = executor

    def execute(self, graph: FederatedGraph, query, ctx=None):
        if ctx is None:
            ctx = {}

        if '_entities' in query.fields_map:
            entities_link = query.fields_map['_entities']
            query = entities_link.node

            queue = Queue(self.executor)
            task_set = queue.fork(None)
            query_workflow = Query(queue, task_set, graph, query, Context(ctx))

            representations = entities_link.options['representations']

            type_ids_map = defaultdict(list)

            for rep in representations:
                typename = rep['__typename']
                keys = graph.extend_node_keys_map[typename]
                # TODO refactor [0], apollo has __resolveReference to pass ident as is
                key = keys[0]
                ident = rep[key]

                type_ids_map[typename].append(ident)

            # TODO hiku federation can support multiple __typename in one query
            #  but I do not know if its required
            for typename in type_ids_map:
                ids = type_ids_map[typename]
                node = graph.nodes_map[typename]
                query_workflow.process_node(node, query, ids)
        else:
            query = InitOptions(graph).visit(query)
            queue = Queue(self.executor)
            task_set = queue.fork(None)
            query_workflow = Query(queue, task_set, graph, query, Context(ctx))

            query_workflow.start()

        res = self.executor.process(queue, query_workflow)
        return res

@@ -28,16 +28,40 @@
 aultdict
+%0Afrom typing import List
 %0A%0Afrom f
@@ -208,16 +208,357 @@
 Queue%0A%0A%0A
+class KeyNotFound(Exception):%0A    def __init__(self, key, rep):%0A        super().__init__(f'No key %22%7Bkey%7D%22 in representation %22%7Brep%7D%22')%0A%0A%0Adef get_keys(graph, typename) -%3E List%5Bint%5D:%0A    node = graph.nodes_map%5Btypename%5D%0A    return %5B%0A        d.args_map%5B'fields'%5D.value for d in%0A        filter(lambda d: d.name == 'key', node.directives)%0A    %5D%0A%0A%0A
 class En
@@ -1275,42 +1275,31 @@
  = g
-raph.extend_node_keys_map%5B
+et_keys(graph, 
 typename
 %5D%0A  
@@ -1286,33 +1286,33 @@
 (graph, typename
-%5D
+)
 %0A               
@@ -1413,16 +1413,99 @@
 keys%5B0%5D%0A
+                if key not in rep:%0A                    raise KeyNotFound(key, rep)%0A
         
@@ -1588,137 +1588,8 @@
 t)%0A%0A
-            # TODO hiku federation can support multiple __typename in one query%0A            #  but I do not know if its required%0A
     

resolve representations using directives