from type_def import *
from helper import *

def find_subst(ctx: dict[str, Type],
               sub: dict[str, Type],
               a: Type,
               b: Type):
    """
    Find substitution s such that ctx(a) = s(sub(b)).
    Note that variables in a and b are considered independent.
    return s.sub if such s exists (. means function composition).
    return error message if type mismatch.
    """
    # is error
    if isinstance(sub, str):
        return sub

    if isinstance(a, TypeVariable) and a.name in ctx:
        a = ctx[a.name]

    if isinstance(a, BotType):
        return sub

    if isinstance(b, TypeVariable):
        if b.name in sub:
            b = sub[b.name]
        else:
            if len(b.constraints) > 0:
                # we cannot handle BotType correctly
                # but that would not be a problem, as the user cannot get
                # BotType in normal circumstances.
                if isinstance(a, TypeVariable):
                    if len(a.constraints) == 0:
                        return f"{b} cannot take value of an unconstrained variable {a}"
                    diff = [v for v in a.constraints if v not in b.constraints]
                    if len(diff) > 0:
                        over = ', '.join([str(v) for v in diff])
                        return f"{b} cannot take value of {a} as {a} can range over [{over}]"
                else:
                    if a not in b.constraints:
                        return f"{b} cannot take value of {a}"
            sub[b.name] = a
            return sub

    if isinstance(a, TypeVariable):
        if a == b:
            return sub
        else:
            return f"{a} can take values other than {b}"

    if isinstance(b, VirtualClassType):
        if isinstance(a, ClassType):
            parents = [a]
        elif isinstance(a, VirtualClassType):
            parents = [a.base]
        else:
            raise CustomError("cannot substitute non-class-type into virtual class")
        while len(parents) > 0:
            current = parents.pop(0)
            if current == b.base:
                return sub
            else:
                parents += current.parents

    if type(a) == type(b):
        if isinstance(a, ParametricType):
            if len(a.params) != len(b.params):
                return f"{a} != {b}"
            old = sub
            for x, y in zip(a.params, b.params):
                old = find_subst(ctx, old, x, y)
            return old
        elif isinstance(a, ClassType) or isinstance(a, PrimitiveType):
            if a.name == b.name:
                return sub
        elif isinstance(a, VirtualClassType):
            return find_subst(ctx, sub, a.base, b.base)
        else:
            raise Exception()

    return f"{a} != {b}"

def resolve_call(obj,
                 fn: str,
                 args: list[Type],
                 assumptions: dict[str, Type],
                 ctx: Context) -> tuple[Type]:
    # TODO: we may want to return the substitution, for monomorphization...
    f_args = None
    f_result = None
    if obj is not None:
        obj = obj.subst(assumptions)
    if obj is None:
        if fn in ctx.functions:
            f = ctx.functions[fn]
            f_args, f_result = TupleType(f[0]), f[1]
        elif fn in ctx.types:
            c = ctx.types[fn]
            if '__init__' in c.methods:
                f = c.methods['__init__']
                if not isinstance(f[0][0], SelfType) or f[1] is not None:
                    raise CustomError(
                        f'__init__ of {c} should accept self and return None'
                    )
                f_args, f_result = TupleType(f[0][1:]), c
            else:
                f_args, f_result = TupleType([]), c
        else:
            raise CustomError(f"No such function {fn}")
    elif isinstance(obj, PrimitiveType) or isinstance(obj, ClassType):
        if fn in obj.methods:
            f = obj.methods[fn]
            if len(f[0]) == 0 or (not isinstance(f[0][0], SelfType) and \
                    f[0][0] != obj):
                raise CustomError('{f} is not a method of {obj}')
            f_args, f_result = TupleType(f[0][1:]), f[1]
        else:
            raise CustomError(f"No such method {fn} in {obj}")
    elif isinstance(obj, VirtualClassType):
        # TODO: may need to emit special annotation that this is a virtual
        # method call?
        if fn in obj.base.methods:
            f = obj.base.methods[fn]
            if len(f[0]) == 0 or not isinstance(f[0][0], SelfType):
                raise CustomError('{f} is not a method of {obj}')
            f_args, f_result = TupleType(f[0][1:]), f[1]
        else:
            raise CustomError(f"No such method {fn} in {c}")
    elif isinstance(obj, TypeVariable):
        # if not constrained, error. otherwise, try all values, and only allow
        # if the results are the same or if they are the same modulo the
        # substitution.
        # expensive operation, but cache should be applicable
        # in order to cache this, our cache must be able to compare equality
        # modulo variable naming... probably not easy either
        if len(obj.constraints) == 0:
            raise CustomError("no methods for unconstrained object")
        results = [resolve_call(obj, fn, args, assumptions | {obj.name: v}, ctx)
                   for v in obj.constraints]
        for v in results[1:]:
            if v != results[0]:
                break
        else:
            # same result
            return results[0]
        results = [v.inv_subst([(a, obj)])
                   for v, a in zip(results, obj.constraints)]
        for v in results[1:]:
            if v != results[0]:
                break
        else:
            # same result
            return results[0]
        raise CustomError("Divergent type after constraints substitution")

    a = TupleType(args)
    subst = find_subst(assumptions, {}, a, f_args)
    if isinstance(subst, str):
        raise CustomError(f"type check failed: {subst}")
    result = f_result.subst(subst)
    if isinstance(result, SelfType):
        return obj
    else:
        return result