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 * string
(* type id = Id of int [@@unboxed] *)

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

let pp_id ppf (Id (n, x)) = Fmt.pf ppf "%s_%d" x n
(* let pp_id ppf (Id n) = Fmt.pf ppf "_%d" n *)

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
  | Open of id

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;
}


(* env *)

module Env = struct
  type t =
    | Empty
    | Obj of {
        pred : t;
        self : id;
        elems : (string, string) Hashtbl.t;
      }
    | Fun of {
        pred : t;
        args : (string * id) list;
        clos : (id, unit) Hashtbl.t;
      }

  let rec find name = function
    | Empty ->
      raise Not_found

    | Fun { pred; args; clos } ->
      begin match List.assoc name args with
        | id -> id, None
        | exception Not_found ->
          let id, fld = find name pred in
          (* mark id's from pred env as needing capture *)
          Hashtbl.replace clos id ();
          id, fld
      end

    | Obj { pred; self; elems } ->
      begin match Hashtbl.find elems name with
        | elem -> self, Some elem
        | exception Not_found -> find name pred
      end
end

(* utils *)

let seq_r a b = Seq (b, a)

let union xs ys =
  List.sort_uniq (compare : id -> id -> int)
    (List.rev_append ys xs)

(* lower *)

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 "get" 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 "%S not in scope" 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 = Hashtbl.create 32 in
    let ends_with_exp =
      List.fold_left
        (fun _ -> function
           | Ast.Item_exp _ -> true
           | Ast.Item_val (name, _)
           | Ast.Item_obj (name, _)
           | Ast.Item_fun (name, _, _) ->
             if Hashtbl.mem elems name then
               compile_error "multiple definitions of %S" name;
             Hashtbl.add elems name name;
             false)
        false
        items
    in
    if is_scope && not ends_with_exp then
      compile_error "scope does not end in expression";

    let self = new_id "obj" in
    let env = Env.Obj { self; elems; pred = env } in

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

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

           | 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 items) in
             fns, name :: vls, init :: ins, clos

           | Ast.Item_fun (name, args, body) ->
             let lam, clos' = lower_lambda self env args body in
             (name, lam) :: fns, vls, ins, union clos clos')
        ([], [], [], [])
        items
    in

    (* if [is_scope], return the last expr, otherwise return the object (self) *)
    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;
        clos;
      },
      List.fold_left
        (fun a b -> Seq (b, a))
        ret
        inits_r      
    )

  and lower_lambda self env args body =
    let args = List.map (fun a -> a, new_id a) args in
    let clos = Hashtbl.create 32 in
    let env = Env.Fun { args; clos; pred = env } in
    let body = lower_exp env body in

    (* wrap body in let bindings to read from the closure *)
    let body, clos =
      Hashtbl.fold
        (fun id () (ir, clos) ->
           if id = self then
             (* [self] isn't "captured"; it IS the closure! *)
             ir, clos
           else
             Let (id, Open id, ir), id :: clos)
        clos
        (body, [])
    in

    let args = List.map snd args in
    { self; args; body }, clos

  in

  let self = new_id "lib" in
  let elems = Hashtbl.create 128 in
  List.iter (fun (name, _) -> Hashtbl.add elems name name) lib;
  let env = Env.Obj { self; elems; pred = Empty } in
  let args = [] in
  let body = lower_block env modl.items in
  { self; args; body }