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


var = TVar([TInt, TBool])
var2 = TVar([TInt])

foo = TObj('Foo', {
    'foo': TFunc([
        FuncArg('a', var, False),
        FuncArg('b', var2, False)
    ], var2, set([var2]))
}, [var])

v = Visitor()
v.assignments['get_x'] = TFunc([FuncArg('in', var, False)], TInt, set([var]))
v.assignments['Foo'] = TFunc([FuncArg('a', var, False)], foo, set([var]))

prelude = set(v.assignments.keys())
print('-----------')
print('prelude')
for key, value in v.assignments.items():
    print(f'{key}: {value}')
print('-----------')

src = """
# a = Foo(1).foo(1, 2)
# b = Foo(1).foo(1, True)
# c = Foo(True).foo(True, 1)
# d = Foo(True).foo(True, True)

a = Foo(1).foo(1, 2)
c = y.foo(True, 1)
y = Foo(True)
"""

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

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

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


# TODO: Occur check