const ast = @import("ast.zig");
const TypeContext = @import("TypeContext.zig");

pub const TypeMap = struct {
    prev: ?*const @This(),
    scheme: *const ast.Scheme,
    match: *const ast.Match,

    pub const Empty = @This(){
        .prev = null,
        .scheme = &ast.Scheme.Empty,
        .match = &ast.Match.Empty,
    };

    pub fn init(match: *const ast.Match, prev: ?*const @This()) @This() {
        return .{
            .prev = prev,
            .scheme = &match.scheme,
            .match = match,
        };
    }

    pub fn initMap(match: *const ast.Match, tyc: *TypeContext, prev: ?*const @This()) !@This() {
        const mm = if (prev) |tm| try tyc.mapMatch(tm, match) else match;
        return TypeMap.init(mm, prev);
    }

    pub fn mapTVar(self: *const @This(), tv: ast.TVar) ?ast.Type {
        // SLOW
        for (self.scheme.tvars, self.match.tvars) |s, m| {
            switch (s) {
                .TVar => |stv| {
                    if (stv.eq(tv)) {
                        return m.Type;
                    }
                },

                else => {},
            }
        } else {
            return (self.prev orelse return null).mapTVar(tv);
        }
    }

    pub fn mapTNum(self: *const @This(), tnum: ast.TNum) ?ast.NumRef {
        return self.match.mapTNum(tnum) orelse (self.prev orelse return null).mapTNum(tnum);
    }

    pub fn mapUnion(self: *const @This(), base: ast.UnionRef) ?ast.UnionRef {
        return self.match.mapUnion(base) orelse (self.prev orelse return null).mapUnion(base);
        // const be = tc.getEnv(base);
        // if (be.env.*) |env| {
        //     return .{ .base = be.base, .env = env };
        // } else {
        //     var tymap: ?*const @This() = self;
        //     while (tymap) |tm| {
        //         for (tm.scheme.envVars, tm.match.envVars) |sb, mb| {
        //             if (sb.id == be.base.id) {
        //                 return tm.mapEnv(mb, tc);
        //             }
        //         }

        //         tymap = tm.prev;
        //     }

        //     unreachable;
        // }
    }

    pub fn tryGetFunctionByID_(self: *const @This(), uid: ast.Association.ID) union(enum) { Fun: ast.Match.AssocRef.InstPair, Id: ast.Association.ID } {
        for (self.scheme.associations, self.match.assocs) |a, r| {
            if (a.uid == uid) {
                return switch (r.*.?) {
                    .Id => |refuid| return self.tryGetFunctionByID_(refuid),
                    .InstFun => |instfun| .{ .Fun = instfun },
                };
            }
        } else {
            return (self.prev orelse return .{ .Id = uid }).tryGetFunctionByID_(uid);
        }
    }
    pub fn tryGetFunctionByID(self: *const @This(), uid: ast.Association.ID) ?ast.Match.AssocRef.InstPair {
        return switch (self.tryGetFunctionByID_(uid)) {
            .Fun => |ip| ip,
            .Id => null,
        };
    }
    pub fn tryGetFunctionOrIDByID(self: *const @This(), uid: ast.Association.ID) ?*?ast.Match.AssocRef {
        for (self.scheme.associations, self.match.assocs) |a, r| {
            if (a.uid == uid) {
                if (r.*) |rr| {
                    return switch (rr) {
                        .Id => |refuid| return self.tryGetFunctionOrIDByID(refuid),
                        .InstFun => r,
                    };
                } else {
                    return r;
                }
            }
        } else {
            return (self.prev orelse return null).tryGetFunctionOrIDByID(uid);
        }
    }
};