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

module Env = struct
  type t =
    | Empty (* TODO: remove me *)
    | Cons of t * t
    | Obj of { self : Code.reg;
               elems : (string, Value.elem) Hashtbl.t }

  let rec find name = function
    | Empty -> raise Not_found
    | Cons (e1, e2) ->
      begin
        try find name e2
        with Not_found -> find name e1
      end
    | Obj { self; elems } ->
      self, Hashtbl.find elems name
end

let compile modl lib =
  let ep = Code.make_block () in
  let currb = ref ep in
  let emit i = Code.extend !currb i in
  let enter b = currb := b in

  let rec compile_exp env rd = function
    | Ast.Literal (Int n) -> emit (LDI (rd, Int n))
    | Ast.Literal True    -> emit (LDI (rd, True))
    | Ast.Literal False   -> emit (LDI (rd, False))
    | Ast.Literal Nil     -> emit (LDI (rd, Nil))

    | Ast.Path path ->
      let obj, loc = compile_path env rd path in
      emit (GET (rd, obj, loc))

    | Ast.Binop (op, e1, e2) ->
      let r1 = rd in
      let r2 = suc rd in
      compile_exp env r1 e1;
      compile_exp env r2 e2;
      begin match op with
        | Ast.Add -> emit (ADD (rd, r1, r2))
        | Ast.Sub -> emit (SUB (rd, r1, r2))
        | Ast.Mul -> emit (MUL (rd, r1, r2))
        | Ast.Div | Ast.Mod -> failwith "Bcc.compile_exp: TODO(Div,Mod)"
        | Ast.Eql -> emit (EQL (rd, r1, r2))
        | Ast.Grt -> emit (GRT (rd, r1, r2))
        | Ast.Lst -> emit (LST (rd, r1, r2))
        | Ast.Not_eql -> emit (EQL (r1, r1, r2)); emit (NOT (rd, r1))
        | Ast.Lst_eql -> emit (GRT (r1, r1, r2)); emit (NOT (rd, r1))
        | Ast.Grt_eql -> emit (LST (r1, r1, r2)); emit (NOT (rd, r1))
      end

    | Ast.Call (fn, args) ->
      let obj, mth = compile_path env rd fn in
      let args =
        List.mapi
          (fun i arg ->
             let ri = off mth (i + 1) in
             compile_exp env ri arg;
             ri)
          args
      in
      emit (CAL (rd, obj, mth, args))

    | Ast.If (e0, e1, e2) ->
      let r0 = rd in
      let b1 = Code.make_block () in
      let b2 = Code.make_block () in
      compile_exp env r0 e0;
      emit (CBR (r0, b1, b2));
      let jp = Code.make_block () in
      enter b1; compile_exp env rd e1; emit (JMP jp);
      enter b2; compile_exp env rd e2; emit (JMP jp);
      enter jp

    | Ast.Fun (_, _) ->
      failwith "Bcc.compile_exp: TODO(Fun)"

    | Ast.Obj items ->
      ignore (compile_block env rd items)

    | Ast.Scope items ->
      begin match compile_block env rd items with
        | Some r -> emit (LDR (rd, r))
        | None -> compile_error "scope does not end with an expression"
      end

  and compile_path env rd path =
    match path with
    | Ast.Var name ->
      let obj, ele =
        try Env.find name env
        with Not_found ->
          compile_error "unbound variable %S" name
      in
      let loc = rd in
      emit (LDI (loc, Value.of_elem ele));
      obj, loc
    | Ast.Ele (lhs, name) ->
      let obj = rd in
      let loc = suc rd in
      compile_exp env obj lhs;
      emit (LOC (loc, obj, name));
      obj, loc

  and compile_block env rd items =
    let elems = Hashtbl.create 100 in
    let n_vals, _, funs_rev =
      List.fold_left
        (fun (nv, nf, fns) -> function
           | Ast.Item_exp _ -> nv, nf, fns
           | Ast.Item_val (name, _)
           | Ast.Item_obj (name, _) ->
             Hashtbl.add elems name (Value.Field nv);
             nv + 1, nf, fns
           | Ast.Item_fun (name, params, body) ->
             Hashtbl.add elems name (Value.Method nf);
             nv, nf + 1, (name, params, body) :: fns)
        (0, 0, [])
        items
    in

    let prevb = !currb in
    let mthds =
      let clo = Code.R 0 in
      let env = Env.Obj { self = clo; elems } in
      List.rev_map
        (fun (_, params, body) ->
           if params <> [] then
             failwith "Bcc.compile_block: TODO(params)";
           let ep = Code.make_block () in
           enter ep;
           let rv = Code.R 1 in
           compile_exp env rv body;
           emit (RET rv);
           Code.Method { n_args = 0; body = { Code.entry = ep } })
        funs_rev
      |> Array.of_list
    in

    enter prevb;
    emit (CON (rd, { n_slots = n_vals; elems; mthds }));

    let r0 = suc rd in
    let r1 = suc r0 in
    let env = Env.Cons (env, Env.Obj { self = rd; elems }) in
    List.fold_left
      (fun _ -> function
         | Ast.Item_exp exp ->
           compile_exp env r0 exp;
           Some r0
         | Ast.Item_val (name, exp) ->
           let el = Hashtbl.find elems name in
           emit (LDI (r0, Value.of_elem el));
           compile_exp env r1 exp;
           emit (SET (r1, rd, r0));
           None
         | Ast.Item_obj (name, body) ->
           (* TODO: it would be ideal if we could CONstruct the empty versions of obj's in
              a sort of "pre-init" phase, before assigning field values. but for now, obj
              items are identical to val's where the rhs is an obj expression. *)
           let el = Hashtbl.find elems name in
           emit (LDI (r0, Value.of_elem el));
           compile_block env r1 body |> ignore;
           emit (SET (r1, rd, r0));
           None
         | Ast.Item_fun (_, _, _) ->
           (* already handled previously *)
           None)
      None
      items
  in

  let init_env =
    let elems = Hashtbl.create 100 in
    List.iteri (fun i (name, _) -> Hashtbl.add elems name (Value.Method i)) lib;
    Env.Obj { self = R 0; elems }
  in
  let rv = Code.R 1 in
  compile_block init_env rv modl.Ast.items |> ignore;
  emit (RET rv);

  { Code.entry = ep }