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 (Code.R i) k = Code.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 sp = function
    | Ir.Var id ->
      get_reg id

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

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

    | ir ->
      emit_exp_s sp ir;
      sp

  and emit_exp_s sp = function
    | Ir.Lit im ->
      emit (LDI (sp, im))

    | Ir.Get path ->
      let obj, loc = emit_path sp path in
      emit (GET (sp, obj, loc))

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

    | Ir.Seq (e1, e2) ->
      emit_exp 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 sp e0 in
      emit (CBR (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 r1 = emit_exp sp e1 in
      emit (match op with
             | Not -> NOT (sp, r1))

    | Ir.Bop (op, e1, e2) ->
      let r1 = emit_exp sp e1 in
      let r2 = emit_exp (suc sp) e2 in
      emit (match op with
             | Add -> ADD (sp, r1, r2)
             | Sub -> SUB (sp, r1, r2)
             | Mul -> MUL (sp, r1, r2)
             | Div -> failwith "Bcc.compile_exp: TODO(Bop(Div))"
             | Mod -> failwith "Bcc.compile_exp: TODO(Bop(Mod))"
             | Eql -> EQL (sp, r1, r2)
             | Grt -> GRT (sp, r1, r2)
             | Lst -> LST (sp, r1, r2))

    | Ir.Call (fn, args) ->
      let obj, mth = emit_path sp fn in
      let args =
        List.mapi
          (fun i arg ->
             let rv = off mth (i + 1) in
             emit_exp_s rv arg; rv)
          args
      in
      emit (CAL (sp, obj, mth, args))

    | 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 n id ->
             Hashtbl.add clos_map id n;
             n + 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 *)
      emit (CON (sp, { n_slots; elems; mthds }));
      Hashtbl.iter
        (fun id idx ->
           let obj = sp in
           let loc = suc sp in
           emit (LDI (loc, Value.of_int idx));
           emit (SET (get_reg id, obj, loc)))
        clos_map

    | Ir.Open id ->
      let idx = 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 (LDI (sp, Value.of_int idx));
      emit (GET (sp, get_reg lam.self, sp))

    | ir ->
      let rv = emit_exp sp ir in
      if rv <> sp then emit (LDR (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

  in

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

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