open Common
open Ast

module Env = struct
  module StringMap = Map.Make (String)

  type 'a t = 'a StringMap.t

  let empty = StringMap.empty
  let add id value env = StringMap.add id value env
  let add_seq seq env = StringMap.add_seq seq env

  let find id env =
    match StringMap.find_opt id env with
    | Some v -> v
    | None -> raise Error.(error (Syntax (UnboundVariable id)))
end

let uniquify t =
  Gensym.init ();
  let rec go (t : 'a t) env : 'a t =
    {
      t with
      desc =
        (match t.desc with
        | Const const -> Const const
        | Var x -> Var (Env.find x env)
        | Lam (args, t') ->
            let args' = List.map Gensym.fresh args in
            let env' =
              Env.add_seq (List.to_seq (List.combine args args')) env
            in
            Lam (args', go t' env')
        | App (t1, terms) -> App (go t1 env, List.map (fun e -> go e env) terms)
        | Prim1 (op, t) -> Prim1 (op, go t env)
        | Prim2 (op, t1, t2) -> Prim2 (op, go t1 env, go t2 env)
        | If (t1, t2, t3) -> If (go t1 env, go t2 env, go t3 env)
        | Let (binding, body) ->
            let name' = Gensym.fresh binding.name in
            let env' = Env.add binding.name name' env in
            Let ({ name = name'; term = go binding.term env }, go body env')
        | Letrec (bindings, body) ->
            let binding_list =
              List.map
                (fun binding -> (binding.name, Gensym.fresh binding.name))
                bindings
            in
            let env' = Env.add_seq (List.to_seq binding_list) env in
            let bindings' =
              List.map
                (fun binding ->
                  {
                    name = Env.find binding.name env';
                    term = go binding.term env';
                  })
                bindings
            in
            Letrec (bindings', go body env')
        | Tuple terms -> Tuple (List.map (fun e -> go e env) terms)
        | Get (t, i) -> Get (go t env, i)
        | Vect (e1, e2) -> Vect (go e1 env, go e2 env)
        | VectGet (e1, e2) -> VectGet (go e1 env, go e2 env)
        | VectSet (e1, e2, e3) -> VectSet (go e1 env, go e2 env, go e3 env)
        | VectLength e -> VectLength (go e env)
        | Ascribe (t, ty) -> Ascribe (go t env, ty)
        | For (x, t1, t2, t3) ->
            let x' = Gensym.fresh x in
            let env' = Env.add x x' env in
            For (x', go t1 env, go t2 env, go t3 env')
        | Variant (x, terms, ty) ->
            let terms' = List.map (fun e -> go e env) terms in
            Variant (x, terms', ty)
        | Match (t, cases, ty) ->
            let cases =
              List.map
                (fun case ->
                  let args = List.map (fun x -> Gensym.fresh x) case.args in
                  let env =
                    Env.add_seq (List.to_seq (List.combine case.args args)) env
                  in
                  { case with args; body = go case.body env })
                cases
            in
            Match (go t env, cases, ty)
        | Seq terms -> Seq (List.map (fun e -> go e env) terms));
    }
  in
  go t Env.empty
