import mem:*
import str:*
import vec:*
import util:*
import tree:*
import ast:*
import binding:*

adt ref_type {
    _unknown,
    _ref,
    _notref
}
fun to_string(r: ref_type): str {
    match (r) {
        ref_type::_unknown() return str("_ref/unknown")
        ref_type::_ref() return str("_ref/ref")
        ref_type::_notref() return str("_ref/notref")
    }
}

adt type {
    _unknown,
    _void,
    _template_placeholder,
    _ptr: *binding<type>,
    _obj: *tree<ast>,
    // triple<pair<vec<pair<is_ref, param_type>>, pair<is_ref, return_type>>, is_variadic, is raw>
    _fun: triple<pair<vec<pair<ref_type, *binding<type>>>, pair<ref_type, *binding<type>>>, bool, bool>,
    _bool,
    _char,
    _uchar,
    _short,
    _ushort,
    _int,
    _uint,
    _long,
    _ulong,
    _float,
    _double
}

fun has_unknown(t: *binding<type>, epoch: binding_epoch): bool {
    match (*t->get_bound_to(epoch)) {
        type::_unknown()    return true
        type::_ptr(p)       return has_unknown(p, epoch)
        type::_obj(o)       return o->data._binding.second.any_true(fun(inner_t: *binding<type>): bool return has_unknown(inner_t, epoch);)
        type::_fun(f)       return has_unknown(f.first.second.second, epoch) || f.first.first.any_true(fun(p: pair<ref_type, *binding<type>>): bool return has_unknown(p.second, epoch);)
    }
    return false
}

fun unify(t1: *binding<type>, t2: *binding<type>, epoch: binding_epoch) {
    println("attempting to unify " + to_string(t1->get_bound_to(epoch)) + " and " + to_string(t2->get_bound_to(epoch)))
    if (is_unknown(t1->get_bound_to(epoch))) {
        t1->set(t2->get_bound_to(epoch), epoch)
    } else if (is_unknown(t2->get_bound_to(epoch))) {
        t2->set(t1->get_bound_to(epoch), epoch)
    } else {
        if (shallow_equality(t1->get_bound_to(epoch), t2->get_bound_to(epoch), epoch)) {
            if (is_fun(t1->get_bound_to(epoch))) {
                unify(t1->get_bound_to(epoch)->_fun.first.second.second, t2->get_bound_to(epoch)->_fun.first.second.second, epoch)
                // unify ref_types
                if (t1->get_bound_to(epoch)->_fun.first.second.first == ref_type::_unknown())
                    t1->get_bound_to(epoch)->_fun.first.second.first = t2->get_bound_to(epoch)->_fun.first.second.first
                if (t2->get_bound_to(epoch)->_fun.first.second.first == ref_type::_unknown())
                    t2->get_bound_to(epoch)->_fun.first.second.first = t1->get_bound_to(epoch)->_fun.first.second.first
                // might be veradic...
                for (var i = 0; i < t1->get_bound_to(epoch)->_fun.first.first.size && i < t2->get_bound_to(epoch)->_fun.first.first.size; i++;) {
                    unify(t1->get_bound_to(epoch)->_fun.first.first[i].second, t2->get_bound_to(epoch)->_fun.first.first[i].second, epoch)
                    if (t1->get_bound_to(epoch)->_fun.first.first[i].first == ref_type::_unknown())
                        t1->get_bound_to(epoch)->_fun.first.first[i].first = t2->get_bound_to(epoch)->_fun.first.first[i].first
                    if (t2->get_bound_to(epoch)->_fun.first.first[i].first == ref_type::_unknown())
                        t2->get_bound_to(epoch)->_fun.first.first[i].first = t1->get_bound_to(epoch)->_fun.first.first[i].first
                }
            } else if (is_ptr(t1->get_bound_to(epoch))) {
                unify(t1->get_bound_to(epoch)->_ptr, t2->get_bound_to(epoch)->_ptr, epoch)
            } else if (is_obj(t1->get_bound_to(epoch))) {
                for (var i = 0; i < t1->get_bound_to(epoch)->_obj->data._binding.second.size; i++;) {
                    unify(t1->get_bound_to(epoch)->_obj->data._binding.second[i], t2->get_bound_to(epoch)->_obj->data._binding.second[i], epoch)
                }
            }
        } else {
            error("Doesn't typecheck! Attempted to unify " + to_string(t1->get_bound_to(epoch)) + " and " + to_string(t2->get_bound_to(epoch)))
        }
    }
}
fun shallow_equality(a: *type, b: *type, epoch: binding_epoch):bool {
    if (is_ptr(a) != is_ptr(b))
        return false
    if (is_ptr(a) && is_ptr(b))
        return true
    match(*a) {
        type::_fun(x) {
            return is_fun(b) && a->_fun.third == b->_fun.third
        }
        type::_obj(x) {
            return is_obj(b) && (get_ast_binding(x, epoch) == get_ast_binding(b->_obj, epoch) || ((!ast_bound(x) || !ast_bound(b->_obj))
                                                                                    && x->data._binding.second.size == b->_obj->data._binding.second.size
                                                                                    && x->data._binding.first == b->_obj->data._binding.first))
        }
    }
    return *a == *b
}

fun inst_temp_type(t: *binding<type>, replacements: ref map<*binding<type>, *binding<type>>, read_epoch: binding_epoch, write_epoch: binding_epoch): *binding<type> {
    match (*t->get_bound_to(read_epoch)) {
        type::_unknown() error("Unknown in temp type")
        type::_obj(o) {
            var binding_types = o->data._binding.second.map(fun(b: *binding<type>): *binding<type> return inst_temp_type(b, replacements, read_epoch, write_epoch);)
            for (var i = 0; i < o->data._binding.second.size; i++;) {
                if (o->data._binding.second[i] != binding_types[i])
                    return binding_p(type::_obj(_binding(o->data._binding.first, binding_types, o->data._binding.third)), write_epoch)
            }
        }
        type::_ptr(p) {
            var cp = inst_temp_type(p, replacements, read_epoch, write_epoch)
            if (cp == p)
                return t
            else
                return binding_p(type::_ptr(cp), write_epoch)
        }
        type::_fun(b) {
            // triple<pair<param_types, return_type>, is_variadic, is raw>
            var rt = make_pair(b.first.second.first, inst_temp_type(b.first.second.second, replacements, read_epoch, write_epoch))
            var pts = b.first.first.map(fun(pt: pair<ref_type, *binding<type>>): pair<ref_type, *binding<type>> return make_pair(pt.first, inst_temp_type(pt.second, replacements, read_epoch, write_epoch));)
            if (rt.second != b.first.second.second)
                return binding_p(type::_fun(make_triple(make_pair(pts, rt), b.second, b.third)), write_epoch)
            for (var i = 0; i < pts.size; i++;)
                if (pts[i].second != b.first.first[i].second)
                    return binding_p(type::_fun(make_triple(make_pair(pts, rt), b.second, b.third)), write_epoch)
            return t
        }
        type::_template_placeholder()  return replacements[t]
    }
    return t
}

fun equality(a: *type, b: *type, count_unknown_as_equal: bool, epoch: binding_epoch): bool {
    /*println("equality of  " + to_string(a) + " and " + to_string(b))*/
    if (count_unknown_as_equal && (is_unknown(a) || is_unknown(b)))
        return true
    match(*a) {
        type::_obj(x) {
            if (!is_obj(b))
                return false
            if (get_ast_binding(x, epoch) == get_ast_binding(b->_obj, epoch))
                return true
            if (!count_unknown_as_equal || (ast_bound(x) && ast_bound(b->_obj)) || x->data._binding.first != b->_obj->data._binding.first ||  x->data._binding.second.size != b->_obj->data._binding.second.size)
                return false
            for (var i = 0; i < x->data._binding.second.size; i++;) {
                if (!equality(x->data._binding.second[i]->get_bound_to(epoch), b->_obj->data._binding.second[i]->get_bound_to(epoch), count_unknown_as_equal, epoch))
                    return false
            }
            return true
        }
        type::_ptr(p) {
            if (!is_ptr(b))
                return false
            return equality(p->get_bound_to(epoch), b->_ptr->get_bound_to(epoch), count_unknown_as_equal, epoch)
        }
        type::_fun(i) {
            if ( !(is_fun(b) && a->_fun.second == b->_fun.second && a->_fun.third == b->_fun.third) )
                return false
            if ( !equality(a->_fun.first.second.second->get_bound_to(epoch), b->_fun.first.second.second->get_bound_to(epoch), count_unknown_as_equal, epoch) )
                return false
            if ( !(a->_fun.first.first.size == b->_fun.first.first.size) )
                return false
            for (var i = 0; i < a->_fun.first.first.size; i++;)
                if ( !equality(a->_fun.first.first[i].second->get_bound_to(epoch), b->_fun.first.first[i].second->get_bound_to(epoch), count_unknown_as_equal, epoch) )
                    return false
            return true
        }
    }
    return *a == *b
}
fun deref_to_string<T>(in: *T): str
    if (in == mem::null<T>())
        return str("null")
    else
        return to_string(in)
fun to_string(it: *type): str {
    match (*it) {
        type::_unknown()               return str("_unknown")
        type::_void()                  return str("_void")
        type::_ptr(p)                  return "*(pre_ref:" + deref_to_string(p->get_bound_to(binding_epoch::pre_ref())) + "/post_ref" + deref_to_string(p->get_bound_to(binding_epoch::post_ref())) + ")"
        type::_obj(b) {
            return "_obj(" + to_string(b->data) + ")"
        }
        type::_fun(b) {
            // triple<pair<param_types, return_type>, is_variadic, is raw>
            var to_ret = str()
            if (b.second)
                to_ret += "_run("
            else
                to_ret += "_fun("
            to_ret += str(", ").join(b.first.first.map(fun(pt: pair<ref_type, *binding<type>>): str return to_string(pt.first) + "(pre_ref:" + deref_to_string(pt.second->get_bound_to(binding_epoch::pre_ref())) + "/post_ref" + deref_to_string(pt.second->get_bound_to(binding_epoch::post_ref())) + ")";))
            if (b.third)
                to_ret += " ..."
            return to_ret + "): " + to_string(b.first.second.first) + ": (pre_ref:" + deref_to_string(b.first.second.second->get_bound_to(binding_epoch::pre_ref())) + "/post_ref:" + deref_to_string(b.first.second.second->get_bound_to(binding_epoch::post_ref())) + ")"
        }
        type::_template_placeholder()  return str("_template_placeholder")
        type::_bool()                  return str("_bool")
        type::_char()                  return str("_char")
        type::_uchar()                 return str("_uchar")
        type::_short()                 return str("_short")
        type::_ushort()                return str("_ushort")
        type::_int()                   return str("_int")
        type::_uint()                  return str("_uint")
        type::_long()                  return str("_long")
        type::_ulong()                 return str("_ulong")
        type::_float()                 return str("_float")
        type::_double()                return str("_double")
    }
    return str("impossible type")
}
fun is_unknown(x: *type): bool {
    match (*x) {
        type::_unknown() return true
    }
    return false
}
fun is_void(x: *type): bool {
    match (*x) {
        type::_void() return true
    }
    return false
}
fun is_ptr(x: *type): bool {
    match (*x) {
        type::_ptr(p) return true
    }
    return false
}
fun is_obj(x: *type): bool {
    match (*x) {
        type::_obj() return true
    }
    return false
}
fun is_fun(x: *type): bool {
    match (*x) {
        type::_fun(b) return true
    }
    return false
}
fun is_template_placeholder(x: *type): bool {
    match (*x) {
        type::_template_placeholder() return true
    }
    return false
}
fun is_bool(x: *type): bool {
    match (*x) {
        type::_bool() return true
    }
    return false
}
fun is_char(x: *type): bool {
    match (*x) {
        type::_char() return true
    }
    return false
}
fun is_uchar(x: *type): bool {
    match (*x) {
        type::_uchar() return true
    }
    return false
}
fun is_short(x: *type): bool {
    match (*x) {
        type::_short() return true
    }
    return false
}
fun is_ushort(x: *type): bool {
    match (*x) {
        type::_ushort() return true
    }
    return false
}
fun is_int(x: *type): bool {
    match (*x) {
        type::_int() return true
    }
    return false
}
fun is_uint(x: *type): bool {
    match (*x) {
        type::_uint() return true
    }
    return false
}
fun is_long(x: *type): bool {
    match (*x) {
        type::_long() return true
    }
    return false
}
fun is_ulong(x: *type): bool {
    match (*x) {
        type::_ulong() return true
    }
    return false
}
fun is_float(x: *type): bool {
    match (*x) {
        type::_float() return true
    }
    return false
}
fun is_double(x: *type): bool {
    match (*x) {
        type::_double() return true
    }
    return false
}
fun is_signed(x: *type): bool {
    match (*x) {
        type::_char()      return true
        type::_int()       return true
        type::_long()      return true
        type::_short()     return true
        type::_float()     return true
        type::_double()    return true

        type::_uchar()     return false
        type::_ushort()    return false
        type::_uint()      return false
        type::_ulong()     return false
    }
    return false
}