import ast
from type_def import *
from top_level import *

test = """
class A:
    a: int
    def foo(a: B):
        pass

class B(A):
    a: str
    def bar(a: list[list[virtual[A]]]) -> A:
        pass
"""

variables = {'X': TypeVariable('X', []), 'Y': TypeVariable('Y', [])}
types = {'int': PrimitiveType('int'), 'str': PrimitiveType('str')}
ctx = Context(variables, types)

ctx, functions, _ = parse_top_level(ctx, ast.parse(test))

for name, t in ctx.types.items():
    if isinstance(t, ClassType):
        print(f"class {t.name}")
        for name, ty in t.fields.items():
            print(f"    {name}: {ty}")
        for name, (args, result, _) in t.methods.items():
            print(f"    {name}: ({', '.join([str(v) for v in args])}) -> {result}")