from __future__ import annotations
import ast
from ast_visitor import Visitor
from nac3_types import *
from primitives import *

src = """
a = 1 // 1

if 1 <= 1 < 2:
    pass

b = test_virtual(virtual(bar, Foo))
b = test_virtual(virtual(foo, Foo))
"""

foo = TObj('Foo', {
    'a': TInt,
}, [])

foo2 = TObj('Foo2', {
    'a': TInt,
}, [])

bar = TObj('Bar', {
    'a': TInt,
    'b': TInt
}, [], [foo])

type_mapping = {
    'Foo': foo,
    'Foo2': foo2,
    'Bar': bar,
}

prelude = {
    'foo': foo,
    'foo2': foo2,
    'bar': bar,
    'test_virtual': TFunc([FuncArg('a', TVirtual(foo), False)], TInt, [])
}

print('-----------')
print('prelude')
for key, value in prelude.items():
    print(f'{key}: {value}')
print('-----------')

v = Visitor(src, prelude.copy(), lambda x: type_mapping[x])

print(src)
v.visit(ast.parse(src))

for a, b in v.virtuals:
    assert isinstance(a, TObj)
    assert b is a or b in a.parents

print('-----------')
print('calls')
for x in v.calls:
    x.check()
    print(f'{x.find()}')

print('-----------')
print('assignments')
for key, value in v.assignments.items():
    if key not in prelude:
        value.check()
        print(f'{key}: {value.find()}')