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

exception Error of string

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

type imm = Value.t
type id = Id of int [@@unboxed]

type uop =
  | Not

and bop =
  | Add
  | Sub
  | Mul
  | Div
  | Mod
  | Eql
  | Grt
  | Lst

type exp =
  | Lit of Value.t
  | Var of id
  | Get of path
  | Set of path * exp
  | Let of id * exp * exp
  | Seq of exp * exp
  | If of exp * exp * exp
  | Uop of uop * exp
  | Bop of bop * exp * exp
  | Call of path * exp list
  | Obj of obj

and path = id * string

and obj = {
  vals : string list;
  funs : (string * lambda) list;
  (* clos : id list; *)
}

and lambda = {
  self : id;
  args : id list;
  body : exp;
}


(* lower *)

let make_id_dispenser () =
  let i = ref (-1) in fun () -> (incr i; Id !i)

module Env = struct
  type t =
    | Empty
    | Cons of t * t
    | Args of (string * id) list
    | Obj of {
        self : id;
        elems : string list;
      }

  let rec find name = function
    | Empty -> raise Not_found
    | Args args ->
      List.assoc name args, None
    | Cons (e1, e2) ->
      begin
        try find name e1 with
          Not_found -> find name e2
      end
    | Obj { self; elems } ->
      if List.mem name elems then
        self, Some name
      else
        raise Not_found
end

let seq_r a b = Seq (b, a)

let lower ~lib (modl : Ast.modl) =
  let new_id = make_id_dispenser () in

  let rec lower_exp env = function
    | Ast.Literal (Int n) -> Lit (Int n)
    | Ast.Literal True -> Lit True
    | Ast.Literal False -> Lit False
    | Ast.Literal Nil  -> Lit Nil

    | Ast.Path path ->
      lower_path env path
        (function
          | `Var id -> Var id
          | `Get (obj, fld) -> Get (obj, fld))

    | Ast.Binop (op, e1, e2) ->
      let not e = Uop (Not, e) in
      let bop, uop = match op with
        | Ast.Add -> Add, Fun.id
        | Ast.Sub -> Sub, Fun.id
        | Ast.Mul -> Mul, Fun.id
        | Ast.Div -> Div, Fun.id
        | Ast.Mod -> Mod, Fun.id
        | Ast.Eql -> Eql, Fun.id
        | Ast.Grt -> Grt, Fun.id
        | Ast.Lst -> Lst, Fun.id
        | Ast.Not_eql -> Eql, not
        | Ast.Grt_eql -> Lst, not
        | Ast.Lst_eql -> Grt, not
      in
      uop (Bop (bop, lower_exp env e1, lower_exp env e2))

    | Ast.Call (fn, args) ->
      lower_path env fn
        (fun fn ->
           let fn_path = match fn with
             | `Var _ -> failwith "Ir.lower_exp: TODO(fcf calls)"
             | `Get (obj, mth) -> (obj, mth)
           in
           let args = List.map (lower_exp env) args in
           Call (fn_path, args))

    | Ast.If (e1, e2, e3) ->
      If (lower_exp env e1, lower_exp env e2, lower_exp env e3)

    | Ast.Obj items ->
      lower_block env items

    | Ast.Scope items ->
      lower_block env items ~is_scope:true

    | Ast.Fun (_, _) ->
      failwith "Ir.lower_exp: TODO(Fun)"

  and lower_path env path k =
    match path with
    | Ast.Ele (obj, fld) ->
      let rhs = lower_exp env obj in
      let lhs = new_id () in
      Let (lhs, rhs, k (`Get (lhs, fld)))

    | Ast.Var name ->
      match Env.find name env with
      | id, None -> k (`Var id)
      | obj, Some fld -> k (`Get (obj, fld))
      | exception Not_found ->
        compile_error "unbound variable %S" name

  and lower_block ?(is_scope = false) env items =
    (* collect names of bindings to form the new environment; also check if a scope ends
       with an expression, if not then it is an error *)
    let elems, ends_with_exp =
      List.fold_left
        (fun (elems, _) -> function
           | Ast.Item_exp _ -> elems, true
           | Ast.Item_val (name, _)
           | Ast.Item_obj (name, _)
           | Ast.Item_fun (name, _, _) -> name :: elems, false)
        ([], false)
        items
    in
    if is_scope && not ends_with_exp then
      compile_error "scope does not end in expression";

    (* build environment for field initializers; NOT for lambda capture *)
    let self = new_id () in
    let env_in = Env.Cons (Obj { self; elems }, env) in

    let funs_r, vals_r, inits_r =
      List.fold_left
        (fun (fns, vls, ins) -> function
           | Ast.Item_exp exp ->
             let init = lower_exp env_in exp in
             fns, vls, init :: ins

           | Ast.Item_val (name, exp) ->
             let init = Set ((self, name), lower_exp env_in exp) in
             fns, name :: vls, init :: ins

           | Ast.Item_obj (name, items) ->
             (* 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 init = Set ((self, name), lower_block env_in items) in
             fns, name :: vls, init :: ins

           | Ast.Item_fun (name, args, body) ->
             let fn = name, compile_lambda env args body in
             fn :: fns, vls, ins)
        ([], [], [])
        items
    in

    (* TODO: closure conversion *)

    (* if [is_scope], return the last expr, otherwise return the object itself *)
    let ret, inits_r = match is_scope, inits_r with
      | true, init :: inits -> init, inits
      | _,    inits         -> Var self, inits
    in
    (* reverse order of inits and decls since they are cons'ed backwards *)
    Let (
      self,
      Obj {
        funs = List.rev funs_r;
        vals = List.rev vals_r;
      },
      List.fold_left
        (fun a b -> Seq (b, a))
        ret
        inits_r      
    )

  and compile_lambda env args body =
    let self = new_id () in
    let args = List.map (fun a -> a, new_id ()) args in
    (* FIXME: environment *)
    let env = ignore env; Env.Args args in
    let body = lower_exp env body in
    let args = List.map snd args in
    { self; args; body }
  in


  let self = new_id () in
  let env =
    (* TODO: lib entries *)
    let _ = lib in
    Env.Empty
  in
  let args = [] in
  let body = lower_block env modl.items in
  { self; args; body }