(*
  This file is part of scilla.

  Copyright (c) 2018 - present Zilliqa Research Pvt. Ltd.
  
  scilla is free software: you can redistribute it and/or modify it under the
  terms of the GNU General Public License as published by the Free Software
  Foundation, either version 3 of the License, or (at your option) any later
  version.
 
  scilla is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 
  You should have received a copy of the GNU General Public License along with
  scilla.  If not, see <http://www.gnu.org/licenses/>.
*)

open Core
open Result.Let_syntax
open Scilla_base
open MonadUtil
open TypeUtil
open ParserUtil
open Literal
open Syntax
module ER = ParserRep
module SR = ParserRep
module SSTypeUtil = TypeUtilities
module SSLiteral = GlobalLiteral
module SSType = SSLiteral.LType
module SSIdentifier = SSType.TIdentifier
module SSName = SSIdentifier.Name
module EvalSyntax = ScillaSyntax (SR) (ER) (SSLiteral)
open SSIdentifier

type ss_field = {
  fname : SSName.t;
  ftyp : SSType.t;
  fval : SSLiteral.t option; (* We may or may not have the value in memory. *)
}

(* The blockchain info is a map from (query_name, query_args) to some info. *)
type bcinfo_state = (string, (string, string) Caml.Hashtbl.t) Caml.Hashtbl.t
type external_state = { caddr : SSLiteral.Bystrx.t; cstate : ss_field list }

type service_mode =
  | IPC of string
  (* Socket address for IPC *)
  | Local

type ss_state =
  | Uninitialized
  | SS of service_mode * ss_field list * external_state list * bcinfo_state

module MakeStateService () = struct
  (* Internal state for the state service. *)
  let ss_cur_state = ref Uninitialized

  (* Sets up the state service object. Should be called before any queries. *)
  let initialize ~sm ~fields ~ext_states ~bcinfo =
    ss_cur_state := SS (sm, fields, ext_states, bcinfo)

  (* Finalize: no more queries. *)
  let finalize () = pure ()

  let assert_init () =
    match !ss_cur_state with
    | Uninitialized -> fail0 ~kind:"StateService: Uninitialized" ?inst:None
    | SS (sm, fields, estates, bcinfo) -> pure (sm, fields, estates, bcinfo)

  let field_type fields fname =
    match
      List.find fields ~f:(fun z -> [%equal: SSName.t] z.fname (get_id fname))
    with
    | Some f -> pure @@ f.ftyp
    | None ->
        fail1
          ~kind:
            (sprintf "StateService: Unable to determine the type of field %s."
               (as_error_string fname))
          ?inst:None
          (ER.get_loc (get_rep fname))

  let fetch_local ~fname ~keys fields =
    let s = fields in
    match
      List.find s ~f:(fun z -> [%equal: SSName.t] z.fname (get_id fname))
    with
    | Some { fname = _; ftyp = MapType _; fval = Some (Map ((kt, vt), mlit)) }
      when not @@ List.is_empty keys ->
        let%bind ret_val_type =
          SSTypeUtil.map_access_type (MapType (kt, vt)) (List.length keys)
        in
        (* Recursively, index with each key and provide the indexed value. *)
        let rec recurser mlit' klist' vt' =
          match klist' with
          | [ k ] ->
              (* Just an assert. *)
              if not @@ [%equal: SSType.t] vt' ret_val_type then
                fail1
                  ~kind:
                    (sprintf
                       "StateService: Failed indexing into map %s. Internal \
                        error."
                       (as_error_string fname))
                  ?inst:None
                  (ER.get_loc (get_rep fname))
              else
                let res = Caml.Hashtbl.find_opt mlit' k in
                pure @@ res
          | k :: krest -> (
              (* we have more nested maps *)
              match Caml.Hashtbl.find_opt mlit' k with
              | Some (SSLiteral.Map ((_, vt''), mlit'')) ->
                  recurser mlit'' krest vt''
              | None ->
                  (* No element found. Return none. *)
                  pure @@ None
              (* The remaining keys cannot be used for indexing as
                 we ran out of nested maps. *)
              | _ ->
                  fail1
                    ~kind:
                      (sprintf
                         "StateService: Cannot index into map %s. Too many \
                          index  keys."
                         (as_error_string fname))
                    ?inst:None
                    (ER.get_loc (get_rep fname)))
          (* this cannot occur. *)
          | [] ->
              fail1
                ~kind:
                  (sprintf
                     "StateService: Internal error in retriving from map %s."
                     (as_error_string fname))
                ?inst:None
                (ER.get_loc (get_rep fname))
        in
        recurser mlit keys vt
    | Some { fname = _; ftyp = _; fval = Some l } -> pure @@ Some l
    | _ ->
        fail1
          ~kind:
            (sprintf "StateService: field \"%s\" not found.\n"
               (as_error_string fname))
          ?inst:None
          (ER.get_loc (get_rep fname))

  let fetch ~fname ~keys =
    let%bind sm, fields, _estates, _bcinfo = assert_init () in
    match sm with
    | IPC socket_addr -> (
        let%bind tp = field_type fields fname in
        let%bind res = StateIPCClient.fetch ~socket_addr ~fname ~keys ~tp in
        if not @@ List.is_empty keys then pure @@ res
        else
          match res with
          | None ->
              fail1
                ~kind:
                  (sprintf "StateService: Field %s not found on IPC server."
                     (as_error_string fname))
                ?inst:None
                (ER.get_loc (get_rep fname))
          | Some _res' -> pure @@ res)
    | Local -> fetch_local ~fname ~keys fields

  let fetch_bcinfo ~query_name ~query_args =
    let%bind sm, _fields, _estates, bcinfo = assert_init () in
    match sm with
    | IPC socket_addr ->
        StateIPCClient.fetch_bcinfo ~socket_addr ~query_name ~query_args
    | Local -> (
        match Caml.Hashtbl.find_opt bcinfo query_name with
        | Some subm -> (
            match Caml.Hashtbl.find_opt subm query_args with
            | Some res -> pure res
            | None ->
                fail0
                  ~kind:
                    ("fetch_bcinfo: query_args " ^ query_args ^ " for "
                   ^ query_name ^ " not found")
                  ?inst:None)
        | None ->
            fail0
              ~kind:("fetch_bcinfo: query_name " ^ query_name ^ " not found")
              ?inst:None)

  (* Common function for external state lookup.
     * If the caddr+fname+keys combination exists:
     *     If ~ignoreval is true: (None, Some type) is returned
     *     if ~ignoreval is false: (Some val, Some type) is returned
     * Else: (None, None) is returned
  *)
  let external_fetch ~caddr ~fname ~keys ~ignoreval =
    let%bind sm, _fields, estates, _bcinfo = assert_init () in
    let caddr_hex = SSLiteral.Bystrx.hex_encoding caddr in
    match sm with
    | IPC socket_addr ->
        StateIPCClient.external_fetch ~socket_addr ~caddr:caddr_hex ~fname ~keys
          ~ignoreval
    | Local -> (
        match
          List.find_map estates ~f:(fun estate ->
              if SSLiteral.Bystrx.equal caddr estate.caddr then
                Some estate.cstate
              else None)
        with
        | Some fields -> (
            match
              List.find_map fields ~f:(fun field ->
                  if SSName.equal field.fname (get_id fname) then
                    Some field.ftyp
                  else None)
            with
            | Some stored_tp ->
                let%bind res = fetch_local ~fname ~keys fields in
                pure (res, Option.map res ~f:(fun _ -> stored_tp))
            | None -> pure (None, None))
        | None -> pure (None, None))

  let update_local ~fname ~keys vopt fields =
    let s = fields in
    match
      List.find s ~f:(fun z -> [%equal: SSName.t] z.fname (get_id fname))
    with
    | Some { fname = _; ftyp = _; fval = Some (Map ((_, vt), mlit)) }
      when not @@ List.is_empty keys ->
        let rec recurser mlit' klist' vt' =
          match klist' with
          (* we're at the last key, update literal. *)
          | [ k ] -> (
              match vopt with
              | Some v ->
                  Caml.Hashtbl.replace mlit' k v;
                  pure @@ s
              | None ->
                  Caml.Hashtbl.remove mlit' k;
                  pure @@ s)
          | k :: krest -> (
              (* we have more nested maps *)
              match Caml.Hashtbl.find_opt mlit' k with
              | Some (SSLiteral.Map ((_, vt''), mlit'')) ->
                  recurser mlit'' krest vt''
              | None ->
                  if is_some vopt then (
                    (* not a delete operation. *)
                    (* We have more keys remaining, but no entry for "k".
                       So create an empty map for "k" and then proceed. *)
                    let mlit'' = Caml.Hashtbl.create 4 in
                    let%bind kt'', vt'' =
                      match vt' with
                      | MapType (keytype, valtype) -> pure (keytype, valtype)
                      | _ ->
                          fail1
                            ~kind:
                              (sprintf
                                 "StateService: Cannot index into map %s due \
                                  to non-map type"
                                 (as_error_string fname))
                            ?inst:None
                            (ER.get_loc (get_rep fname))
                    in
                    Caml.Hashtbl.replace mlit' k (Map ((kt'', vt''), mlit''));
                    recurser mlit'' krest vt'')
                  else
                    (* No point removing a key that doesn't exist. *)
                    pure @@ s
              (* The remaining keys cannot be used for indexing as
                 we ran out of nested maps. *)
              | _ ->
                  fail1
                    ~kind:
                      (sprintf
                         "StateService: Cannot index into map %s. Too many \
                          index keys."
                         (as_error_string fname))
                    ?inst:None
                    (ER.get_loc (get_rep fname)))
          (* this cannot occur. *)
          | [] ->
              fail1
                ~kind:
                  (sprintf "StateService: Internal error in updating map %s."
                     (as_error_string fname))
                ?inst:None
                (ER.get_loc (get_rep fname))
        in
        recurser mlit keys vt
    | Some { fname = f; ftyp = t; fval = Some _ } -> (
        match vopt with
        | Some fval' ->
            let fields' =
              List.filter fields ~f:(fun f ->
                  not ([%equal: SSName.t] f.fname (get_id fname)))
            in
            pure ({ fname = f; ftyp = t; fval = Some fval' } :: fields')
        | None ->
            fail1
              ~kind:
                (sprintf
                   "StateService: Cannot remove non-map value %s from state"
                   (as_error_string fname))
              ?inst:None
              (ER.get_loc (get_rep fname)))
    | _ ->
        fail1
          ~kind:
            (sprintf "StateService: Field \"%s\" not found.\n"
               (as_error_string fname))
          ?inst:None
          (ER.get_loc (get_rep fname))

  let update ~fname ~keys ~value =
    let%bind sm, fields, estates, bcinfo = assert_init () in
    match sm with
    | IPC socket_addr ->
        let%bind tp = field_type fields fname in
        StateIPCClient.update ~socket_addr ~fname ~keys ~value ~tp
    | Local ->
        let%bind fields' = update_local ~fname ~keys (Some value) fields in
        let _ = ss_cur_state := SS (sm, fields', estates, bcinfo) in
        pure ()

  (* Is a key in a map. keys must be non-empty. *)
  let is_member ~fname ~keys =
    let%bind sm, fields, _estates, _bcinfo = assert_init () in
    match sm with
    | IPC socket_addr ->
        let%bind tp = field_type fields fname in
        let%bind res = StateIPCClient.is_member ~socket_addr ~fname ~keys ~tp in
        pure @@ res
    | Local ->
        let%bind v = fetch_local ~fname ~keys fields in
        pure @@ Option.is_some v

  (* Remove a key from a map. keys must be non-empty. *)
  let remove ~fname ~keys =
    let%bind sm, fields, _estates, _bcinfo = assert_init () in
    match sm with
    | IPC socket_addr ->
        let%bind tp = field_type fields fname in
        StateIPCClient.remove ~socket_addr ~fname ~keys ~tp
    | Local ->
        let%bind _ = update_local ~fname ~keys None fields in
        (* We don't need to update ss_cur_state because only map keys can be removed, and that's stateful. *)
        pure ()

  (* Expensive operation, use with care. *)
  let get_full_state () =
    match !ss_cur_state with
    | Uninitialized -> fail0 ~kind:"StateService: Uninitialized" ?inst:None
    | SS (Local, fl, _estates, _bcstate) ->
        mapM fl ~f:(fun f ->
            match f.fval with
            | None ->
                fail0
                  ~kind:
                    (sprintf "StateService: Field %s's value is not known"
                       (SSName.as_error_string f.fname))
                  ?inst:None
            | Some l -> pure (f.fname, f.ftyp, l))
    | SS (IPC _, fl, _estates, _bcstate) ->
        let%bind sl =
          mapM fl ~f:(fun f ->
              let%bind vopt = fetch ~fname:(mk_loc_id f.fname) ~keys:[] in
              match vopt with
              | Some v -> pure (f.fname, f.ftyp, v)
              | None ->
                  fail0
                    ~kind:
                      (sprintf
                         "StateService: Field %s's value not found on server"
                         (SSName.as_error_string f.fname))
                    ?inst:None)
        in
        pure sl
end

(* module MakeStateService *)

module StateServiceInstance = MakeStateService ()
include StateServiceInstance
