open Value

module Hashtbl = struct
  include Stdlib.Hashtbl

  let contains tbl key = find_opt tbl key |> Option.is_some
end

type env_table = (string, lox_value) Hashtbl.t
type environment = { globals : env_table ref; locals : env_table list }

module Env = struct
  type t = environment

  let make () : t = { globals = ref (Hashtbl.create 0); locals = [] }

  let enter (env : t) : t =
    let tbl = Hashtbl.create 0 in
    { env with locals = tbl :: env.locals }

  let push_frame (env : t) : t = { env with locals = [] } |> enter

  (* let exit (env : t) : t =
     assert (not (List.is_empty env.locals));
     { env with locals = List.tl env.locals } *)

  let define_global (env : t) (name : string) (value : lox_value) =
    Hashtbl.replace !(env.globals) name value

  let get_global (env : t) (name : string) : lox_value option = Hashtbl.find_opt !(env.globals) name

  (* Check if `name` is defined in either locals or globals *)
  let is_defined (env : t) (name : string) : bool =
    let _is_defined tbl = Hashtbl.contains tbl name in
    List.fold_left (fun acc tbl -> acc || _is_defined tbl) false (!(env.globals) :: env.locals)

  let define (env : t) (name : string) (value : lox_value) : bool =
    if List.is_empty env.locals then (
      define_global env name value;
      true)
    else
      let tbl = List.hd env.locals in
      if Hashtbl.contains tbl name then false
      else (
        Hashtbl.add tbl name value;
        true)

  let update (env : t) (name : string) (value : lox_value) =
    let rec _update tbls =
      match tbls with
      | [] -> false
      | tbl :: tail ->
          if Hashtbl.contains tbl name then (
            Hashtbl.replace tbl name value;
            true)
          else _update tail
    in
    if _update env.locals then () else Hashtbl.replace !(env.globals) name value

  let get (env : t) (name : string) : lox_value option =
    let rec _get tbls =
      match tbls with
      | [] -> None
      | tbl :: tail -> (
          match Hashtbl.find_opt tbl name with Some x -> Some x | None -> _get tail)
    in
    let val_opt = _get env.locals in
    match val_opt with Some x -> Some x | None -> Hashtbl.find_opt !(env.globals) name
end