module type Elaborate = sig
  type ev

  val elaborate : Typed_ast.e -> ev Ast.t
end

module Elaborate = struct
  type ev = Type.t

  module Env = Map.Make (String)

  let error msg = Common.Error.error (Elaborate msg)
  let lift desc meta = Ast.{ desc; meta }

  let rec elaborate (program : Typed_ast.program) : Type.t Ast.t =
    let env =
      program.decls
      |> List.fold_left
           (fun env -> function
             | Typed_ast.DType (name, ty) -> Env.add name ty env)
           Env.empty
    in
    elaborate_e program.expr env

  and elaborate_e (e : Typed_ast.e) env : Type.t Ast.t =
    let open Typed_ast in
    match e.desc with
    | Const c -> lift (Ast.Const c) (Const.type_of c)
    | Var x ->
        lift (Ast.Var x)
          (match Env.find_opt x env with
          | Some x -> x
          | None -> raise (error (Format.asprintf "unbound variable '%s'" x)))
    | Lam (args, e) ->
        let env' = Env.add_seq (List.to_seq args) env in
        let xs, tys = List.split args in
        let t = elaborate_e e env' in
        lift (Ast.Lam (xs, t)) Type.(Arrow (tys, t.meta))
    | App (e, exprs) ->
        let t = elaborate_e e env in
        let ts = List.map (fun e -> elaborate_e e env) exprs in
        let ty = Option.get Type.(cod t.meta) in
        lift (Ast.App (t, ts)) ty
    | Prim1 (op1, e) ->
        let t = elaborate_e e env in
        let ty = snd (Type.ty_of_op1 op1) in
        lift (Ast.Prim1 (op1, t)) ty
    | Prim2 (op2, e1, e2) ->
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let _, _, ty = Type.ty_of_op2 op2 in
        lift (Ast.Prim2 (op2, t1, t2)) ty
    | If (e1, e2, e3) ->
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let t3 = elaborate_e e3 env in
        lift (Ast.If (t1, t2, t3)) t2.meta
    | Let (binding, body) ->
        let binding' =
          Ast.{ name = binding.name; term = elaborate_e binding.e env }
        in
        let env' = Env.add binding.name binding.ty env in
        let body' = elaborate_e body env' in
        lift (Ast.Let (binding', body')) body'.meta
    | Letrec (bindings, body) ->
        let env' =
          List.fold_left
            (fun env binding -> Env.add binding.name binding.ty env)
            env bindings
        in
        let bindings' =
          List.map
            (fun (binding : binding) ->
              Ast.{ name = binding.name; term = elaborate_e binding.e env' })
            bindings
        in
        let body' = elaborate_e body env' in
        lift (Ast.Letrec (bindings', body')) body'.meta
    | Tuple exprs ->
        let terms = List.map (fun e -> elaborate_e e env) exprs in
        let types = List.map (fun (t : 'a Ast.t) -> t.meta) terms in
        lift (Ast.Tuple terms) (Type.Product types)
    | Get (e, i) ->
        let t = elaborate_e e env in
        let ty =
          match t.meta with
          | Type.Product tys -> List.nth tys i
          | Type.Any -> Type.Any
          | _ -> raise (error "expected tuple")
        in
        lift (Ast.Get (t, i)) ty
    | Vect (e1, e2) ->
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let ty = Type.Vec t2.meta in
        lift (Ast.Vect (t1, t2)) ty
    | VectGet (e1, e2) ->
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let ty =
          match t1.meta with
          | Type.Vec ty' -> ty'
          | Type.Any -> Type.Any
          | _ -> raise (error "expected vector")
        in
        lift (Ast.VectGet (t1, t2)) ty
    | VectSet (e1, e2, e3) ->
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let t3 = elaborate_e e3 env in
        lift (Ast.VectSet (t1, t2, t3)) t1.meta
    | VectLength e ->
        let t = elaborate_e e env in
        lift (Ast.VectLength t) Type.Int
    | Ascribe (e, ty) ->
        let t = elaborate_e e env in
        lift (Ast.Ascribe (t, ty)) ty
    | For (x, e1, e2, e3) ->
        let env' = Env.add x Type.Int env in
        let t1 = elaborate_e e1 env in
        let t2 = elaborate_e e2 env in
        let t3 = elaborate_e e3 env' in
        lift (Ast.For (x, t1, t2, t3)) t3.meta
    | Variant (name, exprs, ty) ->
        let terms = List.map (fun e -> elaborate_e e env) exprs in
        lift (Ast.Variant (name, terms, ty)) ty
    | Match (expr, patterns, ty) ->
        let ty' =
          match ty with
          | Type.Rec (name, ty_inner) -> Type.subst_ty name ty ty_inner
          | _ ->
              raise
                (error
                   (Format.asprintf "expected rec variant type, but got '%a'"
                      Type.pp ty))
        in
        let cases =
          match ty' with
          | Type.Variant cases -> cases
          | _ ->
              raise
                (error
                   (Format.asprintf "expected variant type, but got '%a'"
                      Type.pp ty'))
        in
        let term = elaborate_e expr env in
        let term_asc = lift (Ast.Ascribe (term, ty)) ty in
        let elaborate_pattern (pattern : Typed_ast.pattern) =
          let case : Type.case =
            List.find
              (fun (Type.Case (name, _)) -> name = pattern.casename)
              cases
          in
          let (Type.Case (casename, tys)) = case in
          if List.length pattern.args <> List.length tys then
            raise (error "pattern does not match variant type");
          let env' =
            List.combine pattern.args tys
            |> List.fold_left (fun env (arg, ty) -> Env.add arg ty env) env
          in
          let body = elaborate_e pattern.body env' in
          (Ast.{ casename; args = pattern.args; body }, body.meta)
        in
        let patterns, tys = List.split (List.map elaborate_pattern patterns) in
        lift (Ast.Match (term_asc, patterns, ty)) (List.hd tys)
    | Seq exprs ->
        let terms = List.map (fun e -> elaborate_e e env) exprs in
        let last_term = List.hd (List.rev terms) in
        lift (Ast.Seq terms) last_term.meta
end
