exception UnsupportedAst of string

let raise_unsupported_ast (type a) (f : Format.formatter -> a -> unit) (x : a) =
  let buf = Buffer.create 16 in
  let fmt = Format.formatter_of_buffer buf in
  f fmt x;
  Format.pp_print_flush fmt ();
  raise (UnsupportedAst (Buffer.contents buf))

let raise_unsupported_expr = raise_unsupported_ast Pprintast.expression

type 'a index = { index : int } [@@deriving show { with_path = false }]

let get (arr : 'a Arraylist.t) (i : 'a index) : 'a = Arraylist.get arr i.index

type relop = [ `LT | `LTE | `GT | `GTE | `EQ | `NE ]

let string_of_relop : relop -> string = function
  | `LT -> "<"
  | `LTE -> "<="
  | `GT -> ">"
  | `GTE -> ">="
  | `EQ -> "="
  | `NE -> "<>"

type 'a prim =
  | Add : (expr index * expr index) prim
  | Mul : (expr index * expr index) prim
  | Sub : (expr index * expr index) prim
  | RelOp : (relop * expr index * expr index) prim

and expr =
  | App of expr index * expr index
  | Bool of bool
  | Cons of expr index * expr index
  | If of expr index * expr index * expr index
  | Int of int
  | Lam of string * expr index
  | Let of bool * string * expr index * expr index
  | Nil
  | Prim : 'a prim * 'a -> expr
  | Var of string

type 'a ast = { subexprs : expr Arraylist.t; mutable root : 'a index }

let copy (ast : 'a ast) : 'a ast =
  { subexprs = Arraylist.copy ast.subexprs; root = ast.root }

let get_subexpr (ast : 'a ast) : expr index -> expr = get ast.subexprs

let new_ast () =
  { subexprs = Arraylist.make 0 (Var "<placeholder>"); root = { index = 0 } }

let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
  let add (expr : expr) : expr index =
    let index = Arraylist.length ast.subexprs in
    Arraylist.push ast.subexprs expr;
    { index }
  in
  let binop (prim : (expr index * expr index) prim) : expr index =
    add
      (Lam
         ( "$0",
           add (Lam ("$1", add (Prim (prim, (add (Var "$0"), add (Var "$1"))))))
         ))
  in

  let rec loop (expr : Parsetree.expression) : expr index =
    match expr.pexp_desc with
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "+"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (Add, (loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "-"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (Sub, (loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "*"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (Mul, (loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "<"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`LT, loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "<="; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`LTE, loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident ">"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`GT, loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident ">="; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`GTE, loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "="; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`EQ, loop l, loop r)))
    | Pexp_apply
        ( { pexp_desc = Pexp_ident { txt = Lident "<>"; _ }; _ },
          [ (Nolabel, l); (Nolabel, r) ] ) ->
        add (Prim (RelOp, (`NE, loop l, loop r)))
    | Pexp_apply (f, xs) ->
        List.fold_left
          (fun f -> function
            | Asttypes.Nolabel, x -> add (App (f, loop x))
            | _ -> raise_unsupported_expr expr)
          (loop f) xs
    | Pexp_constant (Pconst_integer (n, None)) -> add (Int (int_of_string n))
    | Pexp_construct
        ({ txt = Lident "::"; _ }, Some { pexp_desc = Pexp_tuple [ hd; tl ]; _ })
      ->
        add (Cons (loop hd, loop tl))
    | Pexp_construct ({ txt = Lident "[]"; _ }, None) -> add Nil
    | Pexp_construct ({ txt = Lident "false"; _ }, None) -> add (Bool false)
    | Pexp_construct ({ txt = Lident "true"; _ }, None) -> add (Bool true)
    | Pexp_fun (Nolabel, None, arg_pat, body) -> (
        match arg_pat.ppat_desc with
        | Ppat_var { txt = name; _ } -> add (Lam (name, loop body))
        | _ -> raise_unsupported_expr expr)
    | Pexp_ident { txt = Lident "+"; _ } -> binop Add
    | Pexp_ident { txt = Lident "-"; _ } -> binop Sub
    | Pexp_ident { txt = Lident "*"; _ } -> binop Mul
    | Pexp_ident { txt = Lident name; _ } -> add (Var name)
    | Pexp_ifthenelse (cond, then_, Some else_) ->
        add (If (loop cond, loop then_, loop else_))
    | Pexp_let
        ( recursive,
          [
            {
              pvb_pat = { ppat_desc = Ppat_var { txt = name; _ }; _ };
              pvb_expr;
              _;
            };
          ],
          body ) ->
        let recursive =
          match recursive with Nonrecursive -> false | Recursive -> true
        in
        add (Let (recursive, name, loop pvb_expr, loop body))
    | _ -> raise_unsupported_expr expr
  in
  loop

let expr_of_parsetree (expr : Parsetree.expression) : expr ast =
  let ast = new_ast () in
  ast.root <- add_expr_to_ast ast expr;
  ast

module Wrap = struct
  let expression (pexp_desc : Parsetree.expression_desc) : Parsetree.expression
      =
    {
      pexp_desc;
      pexp_loc = Location.none;
      pexp_loc_stack = [];
      pexp_attributes = [];
    }

  let ident (name : string) : Longident.t Asttypes.loc =
    { txt = Lident name; loc = Location.none }

  let pattern (ppat_desc : Parsetree.pattern_desc) : Parsetree.pattern =
    {
      ppat_desc;
      ppat_loc = Location.none;
      ppat_loc_stack = [];
      ppat_attributes = [];
    }

  let var (name : string) : string Asttypes.loc =
    { txt = name; loc = Location.none }
end

let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
  let subexpr = get_subexpr ast in

  let list_of_apps expr =
    let rec loop = function
      | App (f, x) ->
          let f, xs = loop (subexpr f) in
          (f, x :: xs)
      | f -> (f, [])
    in
    let f, xs = loop expr in
    (f, List.rev xs)
  in

  let rec loop : expr -> Parsetree.expression = function
    | App _ as expr ->
        let f, xs = list_of_apps expr in
        let xs = List.map (fun x -> (Asttypes.Nolabel, loop (subexpr x))) xs in
        Wrap.expression (Pexp_apply (loop f, xs))
    | Bool b ->
        Wrap.expression (Pexp_construct (Wrap.ident (string_of_bool b), None))
    | Cons (hd, tl) ->
        Wrap.expression
          (Pexp_construct
             ( Wrap.ident "::",
               Some
                 (Wrap.expression
                    (Pexp_tuple [ loop (subexpr hd); loop (subexpr tl) ])) ))
    | If (cond, then_, else_) ->
        Wrap.expression
          (Pexp_ifthenelse
             ( loop (subexpr cond),
               loop (subexpr then_),
               Some (loop (subexpr else_)) ))
    | Int n ->
        Wrap.expression (Pexp_constant (Pconst_integer (Int.to_string n, None)))
    | Lam (x, b) ->
        Wrap.expression
          (Pexp_fun
             ( Nolabel,
               None,
               Wrap.pattern (Ppat_var (Wrap.var x)),
               loop (subexpr b) ))
    | Let (recursive, name, bound, body) ->
        Wrap.expression
          (Pexp_let
             ( (if recursive then Recursive else Nonrecursive),
               [
                 {
                   pvb_pat = Wrap.pattern (Ppat_var (Wrap.var name));
                   pvb_expr = loop (subexpr bound);
                   pvb_constraint = None;
                   pvb_attributes = [];
                   pvb_loc = Location.none;
                 };
               ],
               loop (subexpr body) ))
    | Nil -> Wrap.expression (Pexp_construct (Wrap.ident "[]", None))
    | Prim (Add, (l, r)) ->
        Wrap.expression
          (Pexp_apply
             ( Wrap.expression (Pexp_ident (Wrap.ident "+")),
               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
    | Prim (Sub, (l, r)) ->
        Wrap.expression
          (Pexp_apply
             ( Wrap.expression (Pexp_ident (Wrap.ident "-")),
               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
    | Prim (Mul, (l, r)) ->
        Wrap.expression
          (Pexp_apply
             ( Wrap.expression (Pexp_ident (Wrap.ident "*")),
               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
    | Prim (RelOp, (op, l, r)) ->
        Wrap.expression
          (Pexp_apply
             ( Wrap.expression (Pexp_ident (Wrap.ident (string_of_relop op))),
               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
    | Var name -> Wrap.expression (Pexp_ident (Wrap.ident name))
  in
  loop

let parsetree_of_expr (ast : expr ast) : Parsetree.expression =
  parsetree_of_subexpr ast (get_subexpr ast ast.root)

let show_expr (expr : expr ast) : string =
  let buf = Buffer.create 16 in
  let fmt = Format.formatter_of_buffer buf in
  Pprintast.expression fmt (parsetree_of_expr expr);
  Format.pp_print_flush fmt ();
  Buffer.contents buf

type binder_info = [ `Binder of string | `Bound of expr index ]
type binders = binder_info option array

let get_binders (ast : expr ast) : binders =
  let out = Array.make (Arraylist.length ast.subexprs) None in
  let rec loop (env : (string * expr index) list) (i : expr index) : unit =
    match get_subexpr ast i with
    | App (f, x) ->
        loop env f;
        loop env x
    | Cons (hd, tl) ->
        loop env hd;
        loop env tl
    | If (cond, then_, else_) ->
        loop env cond;
        loop env then_;
        loop env else_
    | Lam (x, b) ->
        out.(i.index) <- Some (`Binder x);
        loop ((x, i) :: env) b
    | Let (false, name, bound, body) ->
        out.(i.index) <- Some (`Binder name);
        loop env bound;
        loop ((name, i) :: env) body
    | Let (true, name, bound, body) ->
        out.(i.index) <- Some (`Binder name);
        loop ((name, i) :: env) bound;
        loop ((name, i) :: env) body
    | Prim (Add, (l, r)) ->
        loop env l;
        loop env r
    | Prim (Sub, (l, r)) ->
        loop env l;
        loop env r
    | Prim (Mul, (l, r)) ->
        loop env l;
        loop env r
    | Prim (RelOp, (_, l, r)) ->
        loop env l;
        loop env r
    | Var name ->
        out.(i.index) <-
          List.assoc_opt name env |> Option.map (fun j -> `Bound j)
    | Bool _ | Int _ | Nil -> ()
  in
  loop [] ast.root;
  out