module Ast = Spice_syntax.Ast
module Code = Spice_runtime.Code
module Value = Spice_runtime.Value

exception Error of string

let compile_error f =
  Fmt.kstr (fun msg -> raise (Error msg)) f

let off (`R i) k = `R (i + k)
let suc r = off r 1

let undef_method =
  Value.Native_function
    (fun _ -> failwith "BUG: method undefined")

let rec compile_lambda ?clos_map (lam : Ir.lambda) =
  let entrypoint = Code.make_block () in
  let currb = ref entrypoint in
  let emit i = Code.extend !currb i in
  let enter b = currb := b in

  let reg_of_id = Hashtbl.create 128 in
  let set_reg id r =
    if Hashtbl.mem reg_of_id id then
      Fmt.failwith "BUG: '%a' reassigned" Ir.pp_id id;
    Hashtbl.add reg_of_id id r
  in
  let get_reg id =
    try Hashtbl.find reg_of_id id with
      Not_found -> Fmt.failwith "BUG: '%a' unassigned" Ir.pp_id id
  in

  let rec emit_exp_v sp = function
    | Ir.Lit v ->
      `Cst v

    | Ir.Var id ->
      (get_reg id :> Code.arg)

    | Ir.Let (id, rhs, bdy) ->
      emit_exp_s sp rhs;
      set_reg id sp;
      emit_exp_v (suc sp) bdy

    | Ir.Seq (e1, e2) ->
      emit_exp_v sp e1 |> ignore;
      emit_exp_v sp e2

    | ir ->
      emit_exp_s sp ir;
      (sp :> Code.arg)

  and emit_exp_s sp : Ir.exp -> unit = function
    | Ir.Get path ->
      let loc = emit_path sp path in
      emit (Get (sp, loc))

    | Ir.Set (path, rhs) ->
      let loc = emit_path sp path in
      let rv = emit_exp_v (suc sp) rhs in
      emit (Set (loc, rv))

    | Ir.Seq (e1, e2) ->
      emit_exp_v sp e1 |> ignore;
      emit_exp_s sp e2

    | Ir.If (e0, e1, e2) ->
      let b1 = Code.make_block () in
      let b2 = Code.make_block () in
      let b3 = Code.make_block () in
      let c = emit_exp_v sp e0 in
      emit (Btr (c, b1, b2));
      enter b1; emit_exp_s sp e1; emit (Jmp b3);
      enter b2; emit_exp_s sp e2; emit (Jmp b3);
      enter b3

    | Ir.Uop (op, e1) ->
      let v1 = emit_exp_v sp e1 in
      let op = match op with Not -> Code.NOT in
      emit (Opr (op, sp, v1))

    | Ir.Bop (op, e1, e2) ->
      let op = match op with
        | Add -> Code.ADD
        | Sub -> Code.SUB
        | Mul -> Code.MUL
        | Div -> Code.DIV
        | Mod -> Code.MOD
        | Eql -> Code.Cmp EQ
        | Grt -> Code.Cmp GT
        | Lst -> Code.Cmp LT
      in
      emit_exp_s sp e1;
      let v2 = emit_exp_v (suc sp) e2 in
      emit (Opr (op, sp, v2))

    | Ir.Call (fn, args) ->
      let fn = emit_path sp fn in
      let args_r, _ =
        List.fold_left
          (fun (args, sp) arg ->
             emit_exp_s sp arg;
             sp :: args, suc sp)
          ([], suc sp)
          args
      in
      emit (Cal (sp, fn, List.rev args_r))

    | Ir.Obj { vals; funs; clos } ->
      (* assign each captured id to a slot *)
      let clos_map = Hashtbl.create 64 in
      let n_slots =
        List.fold_left
          (fun ofs id ->
             Hashtbl.add clos_map id (`Ofs ofs);
             ofs + 1)
          (List.length vals)
          clos
      in

      (* assign each val to a slot *)
      let elems = Hashtbl.create 64 in
      List.iteri
        (fun i name ->
           Hashtbl.add elems name (Value.Field i))
        vals;

      (* compile methods and assign to an index *)
      let mthds = Array.make (List.length funs) undef_method in
      List.iteri
        (fun i (name, lambda) ->
           Hashtbl.add elems name (Value.Method i);
           mthds.(i) <- Code.Function (compile_lambda lambda ~clos_map))
        funs;

      (* construct object and save captured id's *)
      let vtb : Code.vtable = { n_slots; elems; mthds } in
      emit (Con (sp, vtb));
      Hashtbl.iter
        (fun cap_id clos_ofs ->
           let cap_v = (get_reg cap_id :> Code.arg) in
           emit (Set ((sp, clos_ofs), cap_v)))
        clos_map

    | Ir.Open id ->
      let clos = get_reg lam.self in
      let ofs = try Hashtbl.find (Option.get clos_map) id
        with Not_found -> failwith "BUG: %S not captured"
           | Invalid_argument _ -> failwith "BUG: no captured variables"
      in
      emit (Get (sp, (clos, ofs)))

    | ir ->
      let rv = emit_exp_v sp ir in
      if rv <> (sp :> Code.arg) then
        emit (Mov (sp, rv))

  and emit_path sp (obj, fld) =
    let obj = get_reg obj in
    let loc = sp in
    emit (Loc (loc, obj, fld));
    obj, (loc :> Code.ofs)

  in

  (* R0 = self *)
  (* R(i+1) = args[i] *)
  set_reg lam.self (`R 0);
  let sp =
    List.fold_left
      (fun sp arg -> set_reg arg sp; suc sp)
      (`R 1)
      lam.args
  in
  let rv = emit_exp_v sp lam.body in
  emit (Ret rv);

  Code.make_funct
    (List.length lam.args)
    entrypoint