from thinglang.compiler.allocation import LinearMemoryAllocationLayout
from thinglang.compiler.references import ResolvedReference
from thinglang.lexer.tokens.base import LexicalIdentifier
from thinglang.parser.symbols.base import AssignmentOperation
from thinglang.parser.symbols.classes import ThingDefinition, MethodDefinition
from thinglang.parser.symbols.functions import MethodCall
from thinglang.utils.tree_utils import TreeTraversal, inspects


class IndexerContext(object):
    def __init__(self):
        self.instance_members, self.locals, self.current_method, self.current_thing = None, None, None, None
        self.method_index = 0
        self.things = []

    def over(self, thing_definition):
        self.current_thing = thing_definition
        self.things.append(thing_definition)
        self.instance_members = {
            member.name: (idx, member) for idx, member in enumerate(thing_definition.members())
        }

    def set_last_frame_size(self, size):
        if self.current_method:
            self.current_method.frame_size = size

    def set_method(self, node):
        self.current_method = node
        node.index = self.method_index
        self.method_index += 1


class Collator(TreeTraversal):

    def __init__(self, ast):
        super(Collator, self).__init__(ast)
        self.thing_index = itertools.count(0)
        self.method_index = None

    @inspects(ThingDefinition)
    def set_thing_context(self, node: ThingDefinition) -> None:
        node.index = next(self.thing_index)
        node.finalize()
        self.method_index = itertools.count(0)

    @inspects(MethodDefinition)
    def set_method_context(self, node: MethodDefinition) -> None:
        node.index = next(self.method_index)


class Indexer(TreeTraversal):

    def __init__(self, ast):
        super().__init__(ast)
        self.instance_members, self.locals, self.current_method = None, None, None

    def run(self):
        super(Indexer, self).run()
        self.current_method.frame_size = self.locals.next_index

    @inspects(ThingDefinition)
    def set_thing_context(self, node: ThingDefinition) -> None:
        self.instance_members = {
            member.name: (idx, member) for idx, member in enumerate(node.members())
        }

    @inspects(MethodDefinition)
    def set_method_context(self, node: MethodDefinition) -> None:
        if self.current_method:
            self.current_method.frame_size = self.locals.next_index

        self.current_method = node
        self.locals = LinearMemoryAllocationLayout({
            argument: idx for idx, argument in enumerate(node.arguments)
        })

    @inspects(AssignmentOperation)
    def process_assignment_operation(self, node: AssignmentOperation) -> None:
        if node.intent is AssignmentOperation.DECELERATION:
            self.process_declaration(node)
        elif node.intent is AssignmentOperation.REASSIGNMENT:
            self.process_assignment(node)

    @inspects(MethodCall)
    def inspect_method_call(self, node: MethodCall) -> None:
        """
        The reference indexing for method calls involves 2 separate processes - resolving the method target,
        and resolving each and every argument
        """
        node.arguments = [
            arg if arg.STATIC else ResolvedReference(self.locals.get(arg)) for arg in node.arguments
        ]

        if node.target[0].is_self():
            target = self.context.current_thing
            for x in node.target[1:]:
                target = target[x]
            node.resolved_target = ResolvedReference(target.index)

        print(node, node.arguments)

    def process_declaration(self, node):
        if node in self.locals:
            raise Exception('Duplicate decleration')
        else:
            print('Setting a local variable')
            node.target = self.locals.add(node.name)

    def process_assignment(self, node: AssignmentOperation):
        if node.name.implements(LexicalIdentifier):
            node.target = self.locals.get(node.name)
            print('For {}, idx is {}'.format(node.name, self.locals.get(node.name)))
        else:
            raise Exception('cannot resolve references yet!')
@@ -1,28 +1,46 @@
+import itertools%0A%0A
 from thinglang.compiler.allo
@@ -1864,72 +1864,56 @@
 elf.
-instance_members, self.locals, self.current_method = None, None,
+context = IndexerContext()%0A        self.locals =
  Non
@@ -1979,38 +1979,40 @@
       self.c
-urrent_method.
+ontext.set_last_
 frame_size =
@@ -2001,35 +2001,33 @@
 _last_frame_size
- = 
+(
 self.locals.next
@@ -2028,24 +2028,25 @@
 s.next_index
+)
 %0A%0A    @inspe
@@ -2147,122 +2147,26 @@
 elf.
-instance_members = %7B%0A            member.name: (idx, member) for idx, member in enumerate(node.members())%0A        %7D
+context.over(node)
 %0A%0A  
@@ -2273,38 +2273,30 @@
     if self.
-current_method
+locals
 :%0A          
@@ -2299,38 +2299,40 @@
       self.c
-urrent_method.
+ontext.set_last_
 frame_size =
@@ -2329,19 +2329,17 @@
 ame_size
- = 
+(
 self.loc
@@ -2352,16 +2352,17 @@
 xt_index
+)
 %0A%0A      
@@ -2365,37 +2365,41 @@
       self.c
-urren
+ontext.se
 t_method
  = node%0A    
@@ -2378,39 +2378,38 @@
 ntext.set_method
- = 
+(
 node
+)
 %0A        self.lo

Use new context for indexer