module Env = Map.Make (String)

type binding =
  | Var_arg of Ir.arg
  | Var_ele of Ir.arg * Ir.name
  | Var_fun of Ir.arg * Ir.name

let ( @@@ ) f g x = f (g x)
let fcf_mthd_name = "call"

let anf modl =
  let next_uid = ref 0 in
  let gen_id name =
    let uid = !next_uid in
    incr next_uid;
    Ir.Id.{ name; uid }
  in

  let lift_rhs name rhs =
    let tmp = gen_id name in
    let ctx rest = Ir.Let (tmp, rhs, rest) in
    Ir.Arg tmp, ctx
  in

  let rec lower_exp env = function
    | Syn.Literal l -> lift_rhs "lit" (Ir.Literal l)
    | Syn.Path p -> lower_path env p
    | Syn.Call (fn, args) ->
        let (ob_v, ob_ctx), mthd =
          match fn with
          | Syn.Ele (ob, mthd) -> lower_exp env ob, mthd
          | Syn.Var f -> (
              match Env.find f env with
              | Var_fun (ob, mthd) -> (ob, Fun.id), mthd
              | _ -> lower_path env fn, fcf_mthd_name)
        in
        let arg_vs_rev, ctx =
          List.fold_left
            (fun (vs, ctx) arg ->
              let v, v_ctx = lower_exp env arg in
              v :: vs, ctx @@@ v_ctx)
            ([], ob_ctx)
            args
        in
        let fn_e = Ir.Ele (ob_v, mthd) in
        let arg_vs = List.rev arg_vs_rev in
        let v, call_ctx = lift_rhs mthd (Ir.Call (fn_e, arg_vs)) in
        v, ctx @@@ call_ctx
    | Syn.Binop (op, e1, e2) ->
        let v1, v1_ctx = lower_exp env e1 in
        let v2, v2_ctx = lower_exp env e2 in
        let v, op_ctx = lift_rhs "op" (Ir.Binop (op, v1, v2)) in
        v, v1_ctx @@@ v2_ctx @@@ op_ctx
    | Syn.If (e1, e2, e3) ->
        let v1, v1_ctx = lower_exp env e1 in
        let v2, v2_ctx = lower_exp env e2 in
        let v3, v3_ctx = lower_exp env e3 in
        let jn = gen_id "jn" in
        let rv = gen_id "rv" in
        let ctx rest =
          let cont = Ir.{ params = [ rv ]; body = rest } in
          let e2 = v2_ctx (Ir.Jump (jn, [ v2 ])) in
          let e3 = v3_ctx (Ir.Jump (jn, [ v3 ])) in
          v1_ctx (Ir.Let (jn, Cont cont, If (v1, e2, e3)))
        in
        Ir.Arg rv, ctx
    | Syn.Obj items -> lower_obj env items
    (* | Syn.Fun (xs, e) -> () *)
    (* | Syn.Scope items -> () *)
    | _ -> failwith "..."
  and lower_path env = function
    | Syn.Var x -> (
        match Env.find x env with
        | Var_arg v -> v, Fun.id
        | Var_ele (ob, el) -> lift_rhs x (Ir.Get (Ele (ob, el)))
        | Var_fun (ob, el) ->
            (* TODO: special treatment for known FCF's? *)
            lift_rhs x (Ir.Get (Ele (ob, el)))
        | exception Not_found -> failwith (Fmt.str "undefined variable %S" x))
    | Syn.Ele (ob, el) ->
        let ob_v, ob_ctx = lower_exp env ob in
        let el_v, el_ctx = lift_rhs el (Ir.Get (Ele (ob_v, el))) in
        el_v, ob_ctx @@@ el_ctx
  and lower_obj env items =
    (* TODO: detect duplicate item names *)
    let ob = gen_id "ob" in
    let ob_v = Ir.Arg ob in
    (* TODO: get last expression in block *)
    let env, slots, mthds, inits =
      List.fold_left
        (fun (env, slots, mthds, inits) -> function
          | Syn.Item_exp e ->
              let inits = `Exp e :: inits in
              env, slots, mthds, inits
          | Syn.Item_val (name, rhs) ->
              let env = Env.add name (Var_ele (ob_v, name)) env in
              let slots = name :: slots in
              let inits = `Val (name, rhs) :: inits in
              env, slots, mthds, inits
          | Syn.Item_obj (name, body) ->
              let env = Env.add name (Var_ele (ob_v, name)) env in
              let slots = name :: slots in
              let inits = `Obj (name, body) :: inits in
              env, slots, mthds, inits
          | Syn.Item_fun (name, param_names, body) ->
              let env = Env.add name (Var_fun (ob_v, name)) env in
              let mthds = (name, param_names, body) :: mthds in
              env, slots, mthds, inits)
        (env, [], [], [])
        items
    in
    (* mthds was built in reverse order, so rev_map fixes the order *)
    let mthds =
      List.rev_map
        (fun (name, param_names, body) ->
          (* TODO: detect duplicate param names *)
          let params = List.map gen_id param_names in
          let env' =
            List.fold_left2
              (fun env name id -> Env.add name (Var_arg (Arg id)) env)
              env
              param_names
              params
          in
          let body_v, body_ctx = lower_exp env' body in
          let body = body_ctx (Ir.Ret body_v) in
          Ir.{ name; defn = { params; body } })
        mthds
    in
    (* inits was built in reverse order, so make sure to *prepend* contexts
     * when building the final context *)
    let init_ctx =
      List.fold_left
        (fun ctx' -> function
          | `Exp e ->
              let _, ctx = lower_exp env e in
              ctx @@@ ctx'
          | `Val (name, rhs) ->
              let rhs_v, rhs_ctx = lower_exp env rhs in
              let set_ctx rest = Ir.Set (Ele (ob_v, name), rhs_v, rest) in
              rhs_ctx @@@ set_ctx @@@ ctx'
          | `Obj (name, body) ->
              let ob_v, ob_ctx = lower_obj env body in
              let set_ctx rest = Ir.Set (Ele (ob_v, name), ob_v, rest) in
              ob_ctx @@@ set_ctx @@@ ctx')
        Fun.id
        inits
    in
    let slots = List.rev slots in
    let ob_ctx rest = Ir.Let (ob, Obj { mthds; slots }, rest) in
    ob_v, ob_ctx @@@ init_ctx
  in

  let std = gen_id "std" in
  let std_env = Env.empty in
  let std_env =
    List.fold_left
      (fun env fn -> Env.add fn (Var_fun (Arg std, fn)) env)
      std_env
      [ "read"; "print"; "itoa" ]
  in
  let std_env =
    List.fold_left
      (fun env var -> Env.add var (Var_ele (Arg std, var)) env)
      std_env
      [ "fs"; "rand" ]
  in

  let ret_v, ctx = lower_obj std_env modl.Syn.items in
  let params = [ std ] in
  let body = ctx (Ir.Ret ret_v) in
  Ir.{ params; body }