import { isEqual } from "lodash";
import { generate } from "peggy";
import { buildParser, buildParserFile } from "@lezer/generator";
import { Tree, TreeCursor } from "@lezer/common";

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, []);
}

const parserSource = (top: string) => `
  @top result { ${top} }

  @precedence {
    arrow @right,
    app @left
  }

  Expr {
      EUnit { "()" }
    | Var { Ident }
    | Lam { lam "(" Ident ":" Ty ")" "." Expr }
    | App { Expr !app Expr }
    | paren { "(" Expr ")" }
  }

  Ty {
      TUnit { "()" }
    | Arrow { Ty !arrow "->" Ty }
    | paren { "(" Ty ")" }
  }

  @skip { space }

  @tokens {
    space { @whitespace+ }
    lam { "lam" }
    Ident { $[A-Za-z_]+ }
    @precedence { lam, Ident }
  }
`;

const buildApp = (s: string, c: TreeCursor): term => {
  c.firstChild();
  const left = buildExpr(s, c);
  c.nextSibling();
  const right = buildExpr(s, c);
  c.parent();
  return { kind: "app", left, right };
};

const buildLam = (s: string, c: TreeCursor): term => {
  c.firstChild();
  const name = s.slice(c.from, c.to);
  c.nextSibling();
  const ty = buildTy(s, c);
  c.nextSibling();
  const body = buildExpr(s, c);
  c.parent();
  return { kind: "abs", arg: name, ty, body };
};

const buildVar = (s: string, c: TreeCursor): term => {
  c.firstChild();
  const name = s.slice(c.from, c.to);
  c.parent();
  return { kind: "var", name };
};

const buildArrow = (s: string, c: TreeCursor): ty => {
  c.firstChild();
  const left = buildTy(s, c);
  c.nextSibling();
  const right = buildTy(s, c);
  c.parent();
  return { kind: "arr", left, right };
};

const buildTy = (s: string, c: TreeCursor): ty => {
  let result: ty;
  c.firstChild();
  if (c.name === "TUnit") result = { kind: "unit" };
  else if (c.name === "Arrow") result = buildArrow(s, c);
  else if (c.name === "Ty") result = buildTy(s, c);
  else throw new Error("unknown ty " + c.name);
  c.parent();
  return result;
};

const buildExpr = (s: string, c: TreeCursor): term => {
  let result;
  c.firstChild();
  switch (c.name) {
    case "App":
      result = buildApp(s, c);
      break;
    case "Lam":
      result = buildLam(s, c);
      break;
    case "Var":
      result = buildVar(s, c);
      break;
    case "Expr":
      result = buildExpr(s, c);
      break;
    default:
      throw new Error(
        `unknown ${c.name} (${c.from}:${c.to} = "${s.slice(c.from, c.to)}")`,
      );
  }
  c.parent();
  return result;
};

export const parseExpr = (s: string) => {
  const tree = buildParser(parserSource("Expr"), {
    typeScript: true,
    warn: () => {},
  }).parse(s);

  const cursor = tree.cursor();
  cursor.next();
  return buildExpr(s, cursor);
};

export const parseTy = (s: string) => {
  const tree = buildParser(parserSource("Ty"), {
    typeScript: true,
    warn: () => {},
  }).parse(s);

  const cursor = tree.cursor();
  cursor.next();
  return buildTy(s, cursor);
};
// export const parseExpr = (s: string) => parser.parse(s, { startRule: "expr" });
// export const parseTy = (s: string) => parser.parse(s, { startRule: "ty" });