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

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)
        _, _, 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)
except CustomError as e:
    print('Error while type checking:')
    print(e.msg)