module Ast = Spice_syntax.Ast
module Code = Spice_runtime.Code
module Value = Spice_runtime.Value
module Env = Map.Make (String)

type binding = {
  self : Code.regidx;
  slot : Value.slotidx;
}

let compile modl =
  let ep = Code.make_basic_block [] in
  let bb = ref ep in

  let emit is = Code.add_ins !bb is in
  let emit_mov lhs rhs = if rhs <> Code.Reg lhs then emit (MOV (lhs, rhs)) in

  let rec compile_exp env sp = function
    | Ast.Literal Nil -> Code.Cst_nil
    | Ast.Literal True -> Code.Cst_true
    | Ast.Literal False -> Code.Cst_false
    | Ast.Literal (Int i) -> Code.Cst_int i
    | Ast.Path (Var name) -> (
        match Env.find name env with
        | exception Not_found -> Fmt.failwith "unbound: %S" name
        | { self; slot } ->
            emit_mov sp (Code.cst_of_int slot);
            emit (GET (self, sp));
            Reg sp)
    | Ast.Binop (op, lhs, rhs) ->
        let lhs = compile_exp env sp lhs in
        emit_mov sp lhs;
        let rhs = compile_exp env (sp + 1) rhs in
        (match op with
        | Ast.Add -> emit (ADD (sp, rhs))
        | Ast.Sub -> emit (SUB (sp, rhs))
        | Ast.Mul -> emit (MUL (sp, rhs))
        | Ast.Div | Ast.Mod -> failwith "Bcc.compile_exp: TODO: div/mod"
        | Ast.Eql -> emit (EQL (sp, rhs))
        | Ast.Not_eql ->
            emit (EQL (sp, rhs));
            emit (NOT sp)
        | Ast.Lst -> emit (LST (sp, rhs))
        | Ast.Lst_eql ->
            emit (LST (sp, rhs));
            emit (NOT sp)
        | Ast.Grt -> emit (GRT (sp, rhs))
        | Ast.Grt_eql ->
            emit (LST (sp, rhs));
            emit (NOT sp));
        Reg sp
    | Ast.If (cnd, e1, e2) ->
        let l1 = Code.make_basic_block [] in
        let l2 = Code.make_basic_block [] in
        let jp = Code.make_basic_block [] in
        emit (CBR (compile_exp env sp cnd, l1, l2));
        bb := l1;
        emit_mov sp (compile_exp env sp e1);
        emit (JMP jp);
        bb := l2;
        emit_mov sp (compile_exp env sp e2);
        emit (JMP jp);
        bb := jp;
        Reg sp
    | Ast.Obj body -> compile_obj env sp body
    | Ast.Scope body -> compile_scope env sp body
    | _ -> failwith "Bcc.compile_exp: TODO"
  and compile_block env sp items =
    let self = sp in
    let sp = sp + 1 in

    (* construct new env and vtable *)
    let elems = Hashtbl.create (List.length items * 2) in
    let env, n_slots =
      List.fold_left
        (fun (env, n) -> function
          | Ast.Item_fun (_, _, _) | Ast.Item_exp _ -> env, n
          | Ast.Item_obj (name, _) | Ast.Item_val (name, _) ->
              let slot = n in
              let env = Env.add name { self; slot } env in
              Hashtbl.add elems name (Value.Field slot);
              env, n + 1)
        (env, 0)
        items
    in
    let vtable = Value.{ elems; n_slots } in

    (* emit constructor, compile val fields, and get result of final expression *)
    emit (CON (self, vtable));
    let emit_set name rhs =
      let slot = (Env.find name env).slot in
      emit_mov sp (Code.cst_of_int slot);
      emit_mov (sp + 1) rhs;
      emit (SET (self, sp))
    in
    let final_exp =
      List.fold_left
        (fun _ -> function
          | Ast.Item_fun (_, _, _) -> failwith "Bcc: unsupported: methods"
          | Ast.Item_exp exp -> Some (compile_exp env sp exp)
          | Ast.Item_obj (name, body) ->
              emit_set name (compile_obj env (sp + 1) body);
              None
          | Ast.Item_val (name, rhs) ->
              emit_set name (compile_exp env (sp + 1) rhs);
              None)
        None
        items
    in
    self, final_exp
  and compile_obj env sp items =
    let self, _ = compile_block env sp items in
    Code.Reg self
  and compile_scope env sp items =
    let _, final_exp = compile_block env sp items in
    match final_exp with
    | None -> failwith "block must end with an expression"
    | Some ret -> ret
  in

  let env = Env.empty in
  emit_mov 0 (compile_obj env 0 modl.Ast.items);
  emit RET;
  Code.make_program ep