module Ast = Spice_syntax.Ast
module Value = Spice_runtime.Value
open Spice_runtime.Code
open B.Infix

exception Error of string

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

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

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

let rec compile_lambda ?clos_map (lam : Ir.lambda) =
  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 : Ir.exp -> B.t * arg = function
    | Ir.Lit v ->
      B.empty, `Cst v

    | Ir.Var id ->
      B.empty, (get_reg id :> arg)

    | Ir.Let (id, rhs, bdy) ->
      set_reg id sp;
      let bc1 = emit_exp_s sp rhs in
      let bc2, v = emit_exp_v (suc sp) bdy in
      (bc1 +> bc2), v

    | Ir.Seq (e1, e2) ->
      let bc1, _ = emit_exp_v sp e1 in
      let bc2, v = emit_exp_v sp e2 in
      (bc1 +> bc2), v

    | ir ->
      emit_exp_s sp ir, (sp :> arg)

  and emit_exp_s sp : Ir.exp -> B.t = function
    | Ir.Get path ->
      let bc1, loc = emit_path sp path in
      bc1 +> B.get sp loc

    | Ir.Set (path, rhs) ->
      let bc1, loc = emit_path sp path in
      let bc2, rv = emit_exp_v (suc sp) rhs in
      bc1 +> bc2 +> B.set loc rv

    | Ir.Seq (e1, e2) ->
      let bc1, _ = emit_exp_v sp e1 in
      let bc2 = emit_exp_s sp e2 in
      bc1 +> bc2

    | Ir.If (e0, e1, e2) ->
      let bc0, v0 = emit_exp_v sp e0 in
      let bc1 = emit_exp_s sp e1 in
      let bc2 = emit_exp_s sp e2 in
      bc0 +> B.if_ v0 bc1 bc2

    | Ir.Uop (op, e1) ->
      let op = match op with Not -> B.not_ in
      let bc1, v1 = emit_exp_v sp e1 in
      bc1 +> op sp v1

    | Ir.Bop (op, e1, e2) ->
      let op = match op with
        | Add -> B.add
        | Sub -> B.sub
        | Mul -> B.mul
        | Div -> B.div
        | Mod -> B.mod_
        | Eql -> B.ceq
        | Grt -> B.cgt
        | Lst -> B.clt
      in
      let bc1 = emit_exp_s sp e1 in
      let bc2, v2 = emit_exp_v (suc sp) e2 in
      bc1 +> bc2 +> op sp v2

    | Ir.Call (fn, args) ->
      let bc0, fn = emit_path sp fn in
      let argvs = List.mapi (fun i _ -> add sp (i + 1)) args in
      let bc1 = B.concat (List.map2 emit_exp_s argvs args) in
      bc0 +> bc1 +> B.cal sp fn argvs

    | 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) ->
           let funct = compile_lambda lambda ~clos_map in
           Hashtbl.add elems name (Value.Method i);
           mthds.(i) <- Function funct)
        funs;

      (* construct object and save captured id's *)
      let bc0 = B.con sp { n_slots; mthds; elems } in
      (* Hashtbl.iter *)
      (*   clos_map *)
      Hashtbl.fold
        (fun cap_id clos_ofs bc ->
           bc +> B.set (sp, clos_ofs)
                   (get_reg cap_id :> arg))
        clos_map
        bc0

    | 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
      B.get sp (clos, ofs)

    | ir ->
      let bc, rv = emit_exp_v sp ir in
      if rv = (sp :> arg) then
        bc
      else
        bc +> B.mov sp rv

  and emit_path sp (obj, fld) : B.t * loc =
    let obj = get_reg obj in
    let loc = sp in
    B.loc loc obj fld, (obj, (loc :> 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 bc, rv = emit_exp_v sp lam.body in
  let ep = bc |> B.ret rv in

  make_funct
    (List.length lam.args)
    ep