import ast
import sys
import copy
from helper import CustomError
from type_def import SelfType, ClassType
from parse_stmt import parse_stmts
from primitives import simplest_ctx
from top_level import parse_top_level
from inheritance import class_fixup
from lifetime import Lifetime, assign_stmt

if len(sys.argv)  != 2:
    print('please pass the python script name as argument')
    exit()

with open(sys.argv[1], 'r') as f:
    source = f.read()

tree = ast.parse(source, filename=sys.argv[1])

try:
    ctx, fns = parse_top_level(simplest_ctx, tree)
    for c in ctx.types.values():
        if isinstance(c, ClassType):
            class_fixup(c)

    for c, name, fn in fns:
        if c is None:
            params, result, var = ctx.functions[name]
        else:
            params, result, var = ctx.types[c].methods[name]
        # create substitution for type variables
        subst = {k: copy.deepcopy(ctx.variables[k]) for k in var}
        # check if fully annotated all params
        sym_table = {}
        for n, ty in zip(fn.args.args, params):
            if ty is None:
                raise CustomError(
                    'Function parameters must be annotated',
                    fn)
            if isinstance(ty, SelfType):
                ty = ctx.types[c]
            sym_table[n.arg] = ty.subst(subst)
        try:
            print()
            print('checking:')
            print(ast.unparse(fn))
            print('typecheck...')
            _, _, returned = parse_stmts(ctx, sym_table, sym_table, result, fn.body)
            if result is not None and not returned:
                raise CustomError('Function may have no return value', fn)
            print('lifetime check...')
            sym_table = {k: Lifetime(1) for k in sym_table}
            assign_stmt(2, sym_table, fn.body)
            print('OK!')
        except AssertionError:
            pass
except CustomError as e:
    print('Error while type checking:')
    print(e.msg)