import { isEqual } from "lodash";

export type ty = { kind: "arr"; left: ty; right: ty } | { kind: "unit" };

export const unitTy = (): ty => ({ kind: "unit" });

export const arr = (left: ty, right: ty): ty => ({
  kind: "arr",
  left,
  right,
});

export type term =
  | { kind: "unit" }
  | { kind: "var"; name: string }
  | { kind: "app"; left: term; right: term }
  | { kind: "abs"; arg: string; ty: ty; body: term };

export const unit1 = (): term => ({ kind: "unit" });
export const var1 = (name: string): term => ({ kind: "var", name });
export const app1 = (left: term, right: term): term => ({
  kind: "app",
  left,
  right,
});
export const abs1 = (arg: string, ty: ty, body: term): term => ({
  kind: "abs",
  arg,
  ty,
  body,
});

export type term2 =
  | { kind: "unit" }
  | { kind: "var"; idx: number }
  | { kind: "app"; left: term2; right: term2 }
  | { kind: "abs"; ty: ty; body: term2 };

export const unit2 = (): term2 => ({ kind: "unit" });
export const var2 = (idx: number): term2 => ({ kind: "var", idx });
export const app2 = (left: term2, right: term2): term2 => ({
  kind: "app",
  left,
  right,
});
export const abs2 = (ty: ty, body: term2): term2 => ({ kind: "abs", ty, body });

export function convertToDebruijn(t: term): term2 {
  function helper(t: term, ctx: string[]): term2 {
    switch (t.kind) {
      case "unit":
        return { kind: "unit" };
      case "abs":
        return { kind: "abs", ty: t.ty, body: helper(t.body, [t.arg, ...ctx]) };
      case "app":
        return {
          kind: "app",
          left: helper(t.left, ctx),
          right: helper(t.right, ctx),
        };
      case "var":
        return { kind: "var", idx: ctx.indexOf(t.name) };
    }
  }

  return helper(t, []);
}

export function infer(t: term2): ty {
  function check(t: term2, ty: ty, ctx: ty[]): boolean {
    switch (t.kind) {
      case "unit":
        return ty.kind === "unit";
      case "var":
        return isEqual(ctx[t.idx], ty);
      case "app": {
        const fTy = infer(t.left, ctx);
        if (fTy.kind !== "arr") return false;
        const fArgTy = fTy.left;
        const fBodyTy = fTy.right;
        if (!check(t.right, fBodyTy, ctx)) return false;
        return true;
      }
      case "abs":
        // ty MUST be a function
        if (ty.kind !== "arr") return false;
        const { left: tyLeft, right: tyRight } = ty;
        // the types must match
        if (!isEqual(tyLeft, t.ty)) return false;
        return true;
    }
  }

  function infer(t: term2, ctx: ty[]): ty {
    switch (t.kind) {
      case "unit":
        return { kind: "unit" };
      case "abs":
        return {
          kind: "arr",
          left: t.ty,
          right: infer(t.body, [...ctx, t.ty]),
        };
      case "var":
        return ctx[t.idx];
      case "app": {
        const leftTy = infer(t.left, ctx);
        if (leftTy.kind !== "arr") throw new Error("not a function");
        if (!check(t.right, leftTy.right, ctx))
          throw new Error("types don't match");
        // TODO: Ensure that the types match
        return leftTy.right;
      }
    }
  }

  return infer(t, []);
}