import isEqual from "lodash.isequal";
import { match, P } from "ts-pattern";

export type expr =
  | { kind: "var"; x: variable }
  | { kind: "universe"; level: number }
  | { kind: "pi"; abs: abstraction }
  | { kind: "lambda"; abs: abstraction }
  | { kind: "app"; left: expr; right: expr };

export interface abstraction {
  x: variable;
  type: expr;
  body: expr;
}

// =============================================================================
// Substitution

export type variable =
  | { kind: "string"; name: string }
  | { kind: "gensym"; name: string; n: number }
  | { kind: "dummy" };

export const dummy = (): variable => ({ kind: "dummy" });

export interface Substitution {
  x: variable;
  repl: expr;
}

let ctr = 0;
export function refresh(v: variable): variable {
  const n = ctr + 1;
  ctr += 1;
  switch (v.kind) {
    case "gensym":
    case "string": {
      return { kind: "gensym", name: v.name, n };
    }
    case "dummy":
      return { kind: "gensym", name: "_", n };
  }
}

export function subst(s: Substitution[], e: expr): expr {
  switch (e.kind) {
    case "var":
      return s.find((sub) => isEqual(sub.x, e.x))?.repl ?? e;
    case "universe":
      return e;
    case "pi":
      return { kind: "pi", abs: substAbstraction(s, e.abs) };
    case "lambda":
      return { kind: "lambda", abs: substAbstraction(s, e.abs) };
    case "app":
      return { kind: "app", left: subst(s, e.left), right: subst(s, e.right) };
  }
}

export function substAbstraction(
  s: Substitution[],
  { x, type, body }: abstraction,
): abstraction {
  const x2 = refresh(x);
  return {
    x: x,
    type: subst(s, type),
    body: subst([{ x, repl: { kind: "var", x: x2 } }, ...s], body),
  };
}

// =============================================================================
// Type inference

export type context = ContextEntry[];

export interface ContextEntry {
  x: variable;
  type: expr;
  value?: expr | undefined;
}

/** Returns the type of the definition associated with `x` inside the context `ctx` */
export function lookupTy(x: variable, ctx: context): expr | null {
  return ctx.find((entry) => isEqual(entry.x, x))?.type ?? null;
}

/** Returns the value of the definition associated with `x` inside the context `ctx` */
export function lookupValue(x: variable, ctx: context): expr | null {
  return ctx.find((entry) => isEqual(entry.x, x))?.value ?? null;
}

export function extend(
  x: variable,
  t: expr,
  value: expr | undefined,
  ctx: context,
): context {
  return [...ctx, { x: x, type: t, value }];
}

/** Infer the type of the expression `e` in context `ctx` */
export function inferType(ctx: context, e: expr): expr {
  switch (e.kind) {
    case "var": {
      const ty = lookupTy(e.x, ctx);
      if (ty === null) throw new Error(`unknown identifier ${e.x}`);
      return ty;
    }
    case "universe":
      return { kind: "universe", level: e.level + 1 };
    case "pi": {
      const { x, type, body } = e.abs;
      const k1 = inferUniverse(ctx, type);
      const k2 = inferUniverse(extend(x, type, undefined, ctx), e.abs.body);
      return { kind: "universe", level: Math.max(k1, k2) };
    }
    case "lambda": {
      const { x, type, body } = e.abs;
      const te = inferType(extend(x, type, undefined, ctx), body);
      return { kind: "pi", abs: { x, type, body: te } };
    }
    case "app": {
      const { x, type, body } = inferPi(ctx, e.left);
      const te = inferType(ctx, e.right);
      checkEqual(ctx, type, te);
      return subst([{ x, repl: e.right }], body);
    }
  }
}

export function inferUniverse(ctx: context, t: expr): number {
  const u = inferType(ctx, t);
  const n = normalize(ctx, u);
  switch (n.kind) {
    case "universe":
      return n.level;
    default:
      throw new Error("type expected.");
  }
}

export function inferPi(ctx: context, e: expr): abstraction {
  const t = inferType(ctx, e);
  const n = normalize(ctx, t);
  switch (n.kind) {
    case "pi":
      return n.abs;
    default:
      throw new Error("function expected");
  }
}

export function checkEqual(ctx: context, e1: expr, e2: expr) {
  if (!equal(ctx, e1, e2))
    throw new Error(`expressions ${e1} and ${e2} not equal`);
}

// =============================================================================
// Normalization and equality

export function normalize(ctx: context, e: expr): expr {
  switch (e.kind) {
    case "var": {
      const n = lookupTy(e.x, ctx);
      if (!n) throw new Error(`unknown identifier ${e.x}`);
      return n;
    }
    case "app": {
      const e2 = normalize(ctx, e.right);
      const e1 = normalize(ctx, e.left);
      switch (e1.kind) {
        case "lambda":
          return normalize(
            ctx,
            subst([{ x: e1.abs.x, repl: e2 }], e1.abs.body),
          );
        default:
          return { kind: "app", left: e1, right: e2 };
      }
    }
    case "universe":
      return e;
    case "pi":
      return { kind: "pi", abs: normalizeAbstraction(ctx, e.abs) };
    case "lambda":
      return { kind: "lambda", abs: normalizeAbstraction(ctx, e.abs) };
  }
}

export function normalizeAbstraction(
  ctx: context,
  { x, type, body }: abstraction,
): abstraction {
  const t = normalize(ctx, type);
  return { x, type: t, body: normalize(extend(x, t, undefined, ctx), body) };
}

export function equal(ctx: context, e1: expr, e2: expr): boolean {
  function equalHelper(e1: expr, e2: expr): boolean {
    if (e1.kind === "var" && e2.kind === "var") return isEqual(e1.x, e2.x);
    if (e1.kind === "app" && e2.kind === "app")
      return equalHelper(e1.left, e2.left) && equalHelper(e1.right, e2.right);
    if (e1.kind === "universe" && e2.kind === "universe")
      return e1.level === e2.level;
    if (e1.kind === "pi" && e2.kind === "pi")
      return equalAbstraction(e1.abs, e2.abs);
    if (e1.kind === "lambda" && e2.kind === "lambda")
      return equalAbstraction(e1.abs, e2.abs);
    return false;
  }

  function equalAbstraction(
    { x, type: type1, body: body1 }: abstraction,
    { x: y, type: type2, body: body2 }: abstraction,
  ): boolean {
    return (
      equalHelper(type1, type2) &&
      equalHelper(body1, subst([{ x: y, repl: { kind: "var", x } }], body2))
    );
  }

  return equalHelper(normalize(ctx, e1), normalize(ctx, e2));
}