(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2022 Nomadic Labs <contact@nomadic-labs.com>                *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

open Repl_helpers
open Tezos_scoru_wasm

(* Possible step kinds. *)
type eval_step =
  | Tick  (** Tick per tick *)
  | Result  (** Up to Eval (Result (SK_Result _ | SK_Trap _)) *)
  | Kernel_run  (** Up to the end of the current `kernal_run` *)
  | Inbox  (** Until input requested *)

(* Possible commands for the REPL. *)
type commands =
  | Bench
  | Time of commands
  | Show_inbox
  | Show_outbox of int32
  | Show_status
  | Show_durable_storage
  | Show_subkeys of string
  | Show_key of string * printable_value_kind
  | Show_memory of int32 * int * printable_value_kind
  | Dump_function_symbols
  | Step of eval_step
  | Load_inputs
  | Reveal_preimage of string option
  | Reveal_metadata
  | Profile of bool
  | Unknown of string
  | Stop

let parse_eval_step = function
  | "tick" -> Some Tick
  | "result" -> Some Result
  | "kernel_run" -> Some Kernel_run
  | "inbox" -> Some Inbox
  | _ -> None

let parse_memory_commands full_command = function
  | ["at"; address; "for"; length; "bytes"] -> (
      match (Int32.of_string_opt address, int_of_string_opt length) with
      | Some address, Some length -> Show_memory (address, length, `Hex)
      | _, _ -> Unknown full_command)
  | ["at"; address; "for"; length; "bytes"; "as"; "string"] -> (
      match (Int32.of_string_opt address, int_of_string_opt length) with
      | Some address, Some length -> Show_memory (address, length, `String)
      | _, _ -> Unknown full_command)
  | ["at"; address; "as"; kind] -> (
      match
        (Int32.of_string_opt address, integer_value_kind_of_string kind)
      with
      | Some address, Some kind ->
          Show_memory (address, value_kind_length kind, kind)
      | _, _ -> Unknown full_command)
  | _ -> Unknown full_command

let parse_commands s =
  let command = String.split_no_empty ' ' (String.trim s) in
  let rec go = function
    | "time" :: rst ->
        let cmd = go rst in
        Time cmd
    | ["show"; "inbox"] -> Show_inbox
    | ["show"; "outbox"; "at"; "level"; level] -> (
        match Int32.of_string_opt level with
        | Some l -> Show_outbox l
        | None -> Unknown s)
    | ["show"; "status"] -> Show_status
    | ["show"; "durable"; "storage"] -> Show_durable_storage
    | ["show"; "subkeys"; path] | ["ls"; path] -> Show_subkeys path
    | ["show"; "subkeys"] | ["ls"] -> Show_subkeys "/"
    | ["show"; "key"; key] -> Show_key (key, `Hex)
    | ["show"; "key"; key; "as"; kind] -> (
        match printable_value_kind_of_string kind with
        | Some kind -> Show_key (key, kind)
        | None -> Unknown s)
    | "show" :: "memory" :: rest -> parse_memory_commands s rest
    | ["dump"; "function"; "symbols"] -> Dump_function_symbols
    | ["step"; step] -> (
        match parse_eval_step step with Some s -> Step s | None -> Unknown s)
    | ["load"; "inputs"] -> Load_inputs
    | ["reveal"; "preimage"] -> Reveal_preimage None
    | ["reveal"; "preimage"; hex_encoded_preimage] ->
        Reveal_preimage (Some hex_encoded_preimage)
    | ["reveal"; "metadata"] -> Reveal_metadata
    | ["stop"] -> Stop
    | ["bench"] -> Bench
    | ["profile"] -> Profile false
    | ["profile"; "--collapsed"] -> Profile true
    | _ -> Unknown s
  in
  go command

let build_metadata config =
  Tezos_protocol_alpha.Protocol.Alpha_context.Sc_rollup.Metadata.(
    {
      address = config.Config.destination;
      origination_level =
        Tezos_protocol_alpha.Protocol.Alpha_context.Raw_level.of_int32_exn 0l;
    }
    |> Data_encoding.Binary.to_string_exn encoding)

let reveal_preimage_manually retries hash =
  let open Lwt_syntax in
  let* () = Lwt_io.printf "Preimage for hash %s not found.\n%!" hash in
  let rec input_preimage retries : string Lwt.t =
    if retries <= 0 then Stdlib.failwith "Too much retry, aborting"
    else
      let* () = Lwt_io.printf "> " in
      let* input =
        Lwt.catch
          (fun () ->
            let* i = Lwt_io.read_line Lwt_io.stdin in
            return_some i)
          (fun _ -> return_none)
      in
      match Option.bind input (fun bytes -> Hex.to_string (`Hex bytes)) with
      | Some preimage -> Lwt.return preimage
      | None ->
          let* () =
            Lwt_io.printf
              "Error: the preimage is not a valid hexadecimal value.\n%!"
          in
          input_preimage (pred retries)
  in
  input_preimage retries

let reveal_preimage_builtin config retries hash =
  let (`Hex hex) = Hex.of_string hash in
  Lwt.catch
    (fun () ->
      let path = Filename.concat config.Config.preimage_directory hex in
      let ch = open_in path in
      let s = really_input_string ch (in_channel_length ch) in
      close_in ch ;
      Lwt.return s)
    (fun _ -> reveal_preimage_manually retries hex)

let reveals config =
  Tezos_scoru_wasm.Builtins.
    {
      (* We retry 3 times to reveal the preimage manually, this should be
         configurable at some point. *)
      reveal_preimage = reveal_preimage_builtin config 3;
      reveal_metadata = (fun () -> Lwt.return (build_metadata config));
    }

let write_debug =
  Tezos_scoru_wasm.Builtins.Printer (fun msg -> Lwt_io.printf "%s%!" msg)

module Make (Wasm_utils : Wasm_utils_intf.S) = struct
  module Prof = Profiling.Make (Wasm_utils)
  open Wasm_utils

  (* [compute_step tree] is a wrapper around [Wasm_pvm.compute_step] that also
     returns the number of ticks elapsed (whi is always 1). *)
  let compute_step tree =
    let open Lwt_syntax in
    trap_exn (fun () ->
        let+ tree = Wasm.compute_step_with_debug ~write_debug tree in
        (tree, 1L))

  (** [eval_to_result tree] tries to evaluates the PVM until the next `SK_Result`
    or `SK_Trap`, and stops in case of reveal tick or input tick. It has the
    property that the memory hasn't been flushed yet and can be inspected. *)
  let eval_to_result config tree =
    trap_exn (fun () ->
        eval_to_result ~write_debug ~reveal_builtins:(reveals config) tree)

  (* [eval_kernel_run tree] evals up to the end of the current `kernel_run` (or
     starts a new one if already at snapshot point). *)
  let eval_kernel_run config tree =
    let open Lwt_syntax in
    trap_exn (fun () ->
        let* info_before = Wasm.get_info tree in
        let* tree, _ =
          Wasm_fast.compute_step_many
            ~reveal_builtins:(reveals config)
            ~write_debug
            ~stop_at_snapshot:true
            ~max_steps:Int64.max_int
            tree
        in
        let+ info_after = Wasm.get_info tree in
        ( tree,
          Z.to_int64 @@ Z.sub info_after.current_tick info_before.current_tick
        ))

  (* Wrapper around {Wasm_utils.eval_until_input_requested}. *)
  let eval_until_input_requested config tree =
    let open Lwt_syntax in
    trap_exn (fun () ->
        let* info_before = Wasm.get_info tree in
        let* tree =
          eval_until_input_requested
            ~fast_exec:true
            ~reveal_builtins:(Some (reveals config))
            ~write_debug
            ~max_steps:Int64.max_int
            tree
        in
        let+ info_after = Wasm.get_info tree in
        ( tree,
          Z.to_int64 @@ Z.sub info_after.current_tick info_before.current_tick
        ))

  type extra = {functions : string Custom_section.FuncMap.t}

  let produce_flamegraph ~collapse ~max_depth (current_node, remaining_nodes) =
    let filename =
      Time.System.(
        Format.asprintf "wasm-debugger-profiling-%a.out" pp_hum (now ()))
    in
    let path = Filename.(concat (get_temp_dir_name ()) filename) in
    let file = open_out path in
    if remaining_nodes <> [] then
      Format.printf
        "The resulting call graph is inconsistent, please open an issue.\n%!" ;
    Profiling.pp_flamegraph
      ~collapse
      ~max_depth
      Profiling.pp_call
      (Format.formatter_of_out_channel file)
      current_node ;
    close_out file ;
    Format.printf "Profiling result can be found in %s\n%!" path

  let eval_and_profile ~collapse config extra tree =
    let open Lwt_syntax in
    trap_exn (fun () ->
        Format.printf
          "Starting the profiling until new messages are expected. Please note \
           that it will take some time and does not reflect a real computation \
           time.\n\
           %!" ;
        let+ tree, ticks, graph =
          Prof.eval_and_profile
            ~write_debug
            ~reveal_builtins:(reveals config)
            extra.functions
            tree
        in
        produce_flamegraph ~collapse ~max_depth:100 graph ;
        Format.printf "The execution took %a ticks.\n%!" Z.pp_print ticks ;
        tree)

  let set_raw_message_input_step level counter encoded_message tree =
    Wasm.set_input_step (input_info level counter) encoded_message tree

  (* [check_input_request tree] checks the PVM is waiting for inputs, and returns
     an hint to go to the next input request state otherwise. *)
  let check_input_request tree =
    let open Lwt_syntax in
    let* info = Wasm.get_info tree in
    match info.Wasm_pvm_state.input_request with
    | Input_required -> return_ok_unit
    | No_input_required ->
        return_error
          "The PVM is not expecting inputs yet. You can `step inbox` to \
           evaluate the current inbox."
    | Reveal_required _ ->
        (* TODO: https://gitlab.com/tezos/tezos/-/issues/4208 *)
        return_error
          "The PVM is expecting a reveal. This command is not implemented yet."

  (* [load_inputs_gen inboxes level tree] reads the next inbox from [inboxes], set the
     messages at [level] in the [tree], and returns the remaining inbox, the next
     level and the tree. *)
  let load_inputs_gen inboxes level tree =
    let open Lwt_result_syntax in
    match Seq.uncons inboxes with
    | Some (inputs, inboxes) ->
        let* tree =
          trap_exn (fun () ->
              set_full_input_step_gen
                set_raw_message_input_step
                inputs
                level
                tree)
        in
        let*! () =
          Lwt_io.printf
            "Loaded %d inputs at level %ld\n%!"
            (List.length inputs)
            level
        in
        return (tree, inboxes, Int32.succ level)
    | None ->
        let*! () = Lwt_io.printf "No more inputs at level %ld\n%!" level in
        return (tree, inboxes, level)

  let load_inputs inboxes level tree =
    let open Lwt_result_syntax in
    let*! status = check_input_request tree in
    match status with
    | Ok () -> load_inputs_gen inboxes level tree
    | Error msg ->
        Format.printf "%s\n%!" msg ;
        return (tree, inboxes, level)

  (* Eval dispatcher. *)
  let eval level inboxes config step tree =
    let open Lwt_result_syntax in
    let return' ?(inboxes = inboxes) f =
      let* tree, count = f in
      return (tree, count, inboxes, level)
    in
    match step with
    | Tick -> return' (compute_step tree)
    | Result -> return' (eval_to_result config tree)
    | Kernel_run -> return' (eval_kernel_run config tree)
    | Inbox -> (
        let*! status = check_input_request tree in
        match status with
        | Ok () ->
            let* tree, inboxes, level = load_inputs inboxes level tree in
            let* tree, ticks = eval_until_input_requested config tree in
            return (tree, ticks, inboxes, level)
        | Error _ -> return' (eval_until_input_requested config tree))

  let profile ~collapse level inboxes config extra tree =
    let open Lwt_result_syntax in
    let*! status = check_input_request tree in
    match status with
    | Ok () ->
        let* tree, inboxes, level = load_inputs inboxes level tree in
        let* tree = eval_and_profile ~collapse config extra tree in
        return (tree, inboxes, level)
    | Error _ ->
        let* tree = eval_and_profile ~collapse config extra tree in
        return (tree, inboxes, level)

  let pp_input_request ppf = function
    | Wasm_pvm_state.No_input_required -> Format.fprintf ppf "Evaluating"
    | Input_required -> Format.fprintf ppf "Waiting for input"
    | Reveal_required Reveal_metadata ->
        Format.fprintf ppf "Waiting for metadata"
    | Reveal_required (Reveal_raw_data hash) ->
        Format.fprintf ppf "Waiting for reveal: %a" Hex.pp @@ Hex.of_string hash

  (* [show_status tree] show the state of the PVM. *)
  let show_status tree =
    let open Lwt_syntax in
    let* state = Wasm.Internal_for_tests.get_tick_state tree in
    let* info = Wasm.get_info tree in
    Lwt_io.printf
      "%s\n%!"
      (Format.asprintf
         "Status: %a\nInternal_status: %a"
         pp_input_request
         info.Wasm_pvm_state.input_request
         pp_state
         state)

  (* [step level inboxes config kind tree] evals according to the step kind and
     prints the number of ticks elapsed and the new status. *)
  let step level inboxes config kind tree =
    let open Lwt_result_syntax in
    let* tree, ticks, inboxes, level = eval level inboxes config kind tree in
    let*! () = Lwt_io.printf "Evaluation took %Ld ticks so far\n" ticks in
    let*! () = show_status tree in
    return (tree, inboxes, level)

  let bench config tree =
    let open Lwt_syntax in
    let action =
      Octez_smart_rollup_wasm_benchmark_lib.Scenario.exec_slow
        ~reveal_builtins:(reveals config)
    in
    let step =
      Octez_smart_rollup_wasm_benchmark_lib.Scenario.make_scenario_step
        ""
        action
    in
    let eval = Octez_smart_rollup_wasm_benchmark_lib.Scenario.apply_step step in
    let* benchmark, tree = eval tree in
    let* () =
      Lwt_io.printf
        "%s\n%!"
        (Format.asprintf
           "%a"
           Octez_smart_rollup_wasm_benchmark_lib.Data.pp_analysis
           benchmark)
    in
    return tree

  (* [show_inbox tree] prints the current input buffer and the number of messages
     it contains. *)
  let show_inbox tree =
    let open Lwt_syntax in
    let* input_buffer = Wasm.Internal_for_tests.get_input_buffer tree in
    let* messages =
      Tezos_lazy_containers.Lazy_vector.(
        Mutable.ZVector.snapshot input_buffer |> ZVector.to_list)
    in
    let messages_sorted =
      List.sort Messages.compare_input_buffer_messages messages
    in
    let pp_message ppf
        Tezos_webassembly_interpreter.Input_buffer.
          {raw_level; message_counter; payload} =
      Format.fprintf
        ppf
        {|{ raw_level: %ld;
  counter: %s
  payload: %a }|}
        raw_level
        (Z.to_string message_counter)
        Messages.pp_input
        payload
    in
    let pp_messages () =
      Format.asprintf
        "%a"
        (Format.pp_print_list
           ~pp_sep:(fun ppf () -> Format.fprintf ppf "\n")
           pp_message)
        messages_sorted
    in

    let size =
      Tezos_lazy_containers.Lazy_vector.Mutable.ZVector.num_elements
        input_buffer
    in
    Lwt_io.printf
      "Inbox has %s messages:\n%s\n%!"
      (Z.to_string size)
      (pp_messages ())

  (* [show_outbox_gen tree level] prints the outbox messages for a given level. *)
  let show_outbox_gen tree level =
    let open Lwt_syntax in
    let* output_buffer = Wasm.Internal_for_tests.get_output_buffer tree in
    let* level_vector =
      Tezos_webassembly_interpreter.Output_buffer.get_outbox output_buffer level
    in
    let* messages =
      Tezos_lazy_containers.Lazy_vector.(level_vector |> ZVector.to_list)
    in
    let pp_messages () =
      Format.asprintf
        "%a"
        (Format.pp_print_list
           ~pp_sep:(fun ppf () -> Format.fprintf ppf "\n")
           Messages.pp_output)
        messages
    in
    let size =
      Tezos_lazy_containers.Lazy_vector.ZVector.num_elements level_vector
    in
    Lwt_io.printf
      "Outbox has %s messages:\n%s\n%!"
      (Z.to_string size)
      (pp_messages ())

  (* [show_outbox tree level] prints the outbox messages for a given level. *)
  let show_outbox tree level =
    Lwt.catch
      (fun () -> show_outbox_gen tree level)
      (fun _ -> Lwt_io.printf "No outbox found at level %ld\n%!" level)

  (* [find_key_in_durable] retrieves the given [key] from the durable storage in
     the tree. Returns `None` if the key does not exists. *)
  let find_key_in_durable tree key =
    let open Lwt_syntax in
    let* durable = Wasm_utils.wrap_as_durable_storage tree in
    let durable = Tezos_scoru_wasm.Durable.of_storage_exn durable in
    Tezos_scoru_wasm.Durable.find_value durable key

  (* [print_durable ~depth ~show_values ~path tree] prints the keys in the durable
     storage from the given path and their values in their hexadecimal
     representation. By default, it prints from the root of the durable
     storage. *)
  let print_durable ?(depth = 10) ?(show_values = true) ?(path = []) tree =
    let open Lwt_syntax in
    let durable_path = "durable" :: path in
    let* path_exists = Wasm_utils.Ctx.Tree.mem_tree tree durable_path in
    if path_exists then
      Wasm_utils.Ctx.Tree.fold
        ~depth:(`Le depth)
        tree
        ("durable" :: path)
        ~order:`Sorted
        ~init:()
        ~f:(fun key tree () ->
          let full_key = String.concat "/" key in
          (* If we need to show the values, we show every keys, even the root and
             '@'. *)
          if show_values then
            let+ value = Wasm_utils.Ctx.Tree.find tree [] in
            let value = Option.value ~default:(Bytes.create 0) value in
            Format.printf "/%s\n  %a\n%!" full_key Hex.pp (Hex.of_bytes value)
          else if key <> [] && key <> ["@"] then
            return (Format.printf "/%s\n%!" full_key)
          else return_unit)
    else
      Lwt.return
      @@ Format.printf
           "The path /%s is not available in the durable storage\n%!"
           (String.concat "/" path)

  (* [show_durable] prints the durable storage from the tree. *)
  let show_durable tree = print_durable ~depth:10 tree

  (* [show_subkeys tree path] prints the direct subkeys under the given path. *)
  let show_subkeys tree path =
    let invalid_path () =
      Lwt.return @@ Format.printf "Invalid path, it must start with '/'\n%!"
    in
    match String.index_opt path '/' with
    | Some i when i <> 0 -> invalid_path ()
    | None -> invalid_path ()
    | Some _ ->
        print_durable
          ~depth:1
          ~path:(String.split_no_empty '/' path)
          ~show_values:false
          tree

  let show_value kind value =
    let printable = Repl_helpers.print_wasm_encoded_value kind value in
    match printable with
    | Ok s -> s
    | Error err ->
        Format.asprintf
          "Error: %s. Defaulting to hexadecimal value\n%a"
          err
          Hex.pp
          (Hex.of_string value)

  (* [show_key_gen tree key] looks for the given [key] in the durable storage and
     print its value in hexadecimal format. *)
  let show_key_gen tree key kind =
    let open Lwt_syntax in
    let* value = find_key_in_durable tree key in
    match value with
    | None ->
        Format.printf "Key not found\n%!" ;
        return_unit
    | Some v ->
        let+ str_value =
          Tezos_lazy_containers.Chunked_byte_vector.to_string v
        in
        Format.printf "%s\n%!" @@ show_value kind str_value

  (* [show_key tree key] looks for the given [key] in the durable storage and
     print its value in hexadecimal format. Prints errors in case the key is
     invalid or not existing. *)
  let show_key tree key kind =
    Lwt.catch
      (fun () ->
        let key = Tezos_scoru_wasm.Durable.key_of_string_exn key in
        show_key_gen tree key kind)
      (function
        | Tezos_scoru_wasm.Durable.Invalid_key _ ->
            Lwt_io.printf "Invalid key\n%!"
        | Tezos_scoru_wasm.Durable.Value_not_found ->
            Lwt_io.printf "No value found for key\n%!"
        | Tezos_scoru_wasm.Durable.Tree_not_found ->
            Lwt_io.printf "No tree found for key\n%!"
        | exn ->
            Lwt_io.printf "Unknown exception: %s\n%!" (Printexc.to_string exn))

  exception Cannot_inspect_memory of string

  (* [load_memory tree] finds the memory module 0 from the tree, only and only if
     the PVM is in an Init or Eval state. *)
  let load_memory tree =
    let open Lwt_syntax in
    let* state = Wasm.Internal_for_tests.get_tick_state tree in
    let* module_inst =
      match state with
      | Eval _ | Init _ -> Wasm.Internal_for_tests.get_module_instance_exn tree
      | _ -> raise (Cannot_inspect_memory (Format.asprintf "%a" pp_state state))
    in
    Lwt.catch
      (fun () ->
        Tezos_lazy_containers.Lazy_vector.Int32Vector.get
          0l
          module_inst.memories)
      (fun _ ->
        raise Tezos_webassembly_interpreter.Eval.Missing_memory_0_export)

  (* [show_memory tree address length] loads the [length] bytes at address
     [address] in the memory, and prints it in its hexadecimal representation. *)
  let show_memory tree address length kind =
    let open Lwt_syntax in
    Lwt.catch
      (fun () ->
        let* memory = load_memory tree in
        let* value =
          Tezos_webassembly_interpreter.Memory.load_bytes memory address length
        in
        Lwt_io.printf "%s\n%!" @@ show_value kind value)
      (function
        | Cannot_inspect_memory state ->
            Lwt_io.printf
              "Error: Cannot inspect memory during internal state %s\n%!"
              state
        | exn -> Lwt_io.printf "Error: %s\n%!" (Printexc.to_string exn))

  let dump_function_symbols extra =
    let functions =
      Format.asprintf "%a" Custom_section.pp_function_subsection extra.functions
    in
    Lwt_io.printf "Functions:\n%s\n" functions

  (* [reveal_preimage config hex tree] checks the current state is waiting for a
     preimage, parses [hex] as an hexadecimal representation of the data or use
     the builtin if none is given, and does a reveal step. *)
  let reveal_preimage config bytes tree =
    let open Lwt_syntax in
    let* info = Wasm.get_info tree in
    match
      ( info.Tezos_scoru_wasm.Wasm_pvm_state.input_request,
        Option.bind bytes (fun bytes -> Hex.to_bytes (`Hex bytes)) )
    with
    | ( Tezos_scoru_wasm.Wasm_pvm_state.(Reveal_required (Reveal_raw_data _)),
        Some preimage ) ->
        Wasm.reveal_step preimage tree
    | Reveal_required (Reveal_raw_data hash), None ->
        Lwt.catch
          (fun () ->
            let* preimage = reveal_preimage_builtin config 1 hash in
            Wasm.reveal_step (String.to_bytes preimage) tree)
          (fun _ -> return tree)
    | _ ->
        let+ () = Lwt_io.print "Error: Not in a reveal step\n%!" in
        tree

  let reveal_metadata config tree =
    let open Lwt_syntax in
    let* info = Wasm.get_info tree in
    match info.Tezos_scoru_wasm.Wasm_pvm_state.input_request with
    | Tezos_scoru_wasm.Wasm_pvm_state.(Reveal_required Reveal_metadata) ->
        let data = build_metadata config in
        Wasm.reveal_step (Bytes.of_string data) tree
    | _ ->
        let+ () = Lwt_io.print "Error: Not in a reveal step\n%!" in
        tree

  (* [handle_command command tree inboxes level] dispatches the commands to their
     actual implementation. *)
  let handle_command c config extra tree inboxes level =
    let open Lwt_result_syntax in
    let command = parse_commands c in
    let return ?(tree = tree) ?(inboxes = inboxes) () =
      return (tree, inboxes, level)
    in
    let rec go = function
      | Bench ->
          (* TODO: https://gitlab.com/tezos/tezos/-/issues/5615
             let*! tree = bench config tree in
             return ~tree ()
          *)
          let*! () =
            Lwt_io.printf
              "This command is temporarily unavailable.\n\
               See https://gitlab.com/tezos/tezos/-/issues/5615\n\
               %!"
          in
          return ()
      | Time cmd ->
          let t = Time.System.now () in
          let* res = go cmd in
          let t' = Time.System.now () in
          let*! () =
            Lwt_io.printf
              "took %s\n%!"
              (Format.asprintf "%a" Ptime.Span.pp (Ptime.diff t' t))
          in
          Lwt_result_syntax.return res
      | Load_inputs -> load_inputs inboxes level tree
      | Show_status ->
          let*! () = show_status tree in
          return ()
      | Step kind -> step level inboxes config kind tree
      | Show_inbox ->
          let*! () = show_inbox tree in
          return ()
      | Show_outbox level ->
          let*! () = show_outbox tree level in
          return ()
      | Show_durable_storage ->
          let*! () = show_durable tree in
          return ()
      | Show_subkeys path ->
          let*! () = show_subkeys tree path in
          return ()
      | Show_key (key, kind) ->
          let*! () = show_key tree key kind in
          return ()
      | Show_memory (address, length, kind) ->
          let*! () = show_memory tree address length kind in
          return ()
      | Dump_function_symbols ->
          let*! () = dump_function_symbols extra in
          return ()
      | Reveal_preimage bytes ->
          let*! tree = reveal_preimage config bytes tree in
          return ~tree ()
      | Reveal_metadata ->
          let*! tree = reveal_metadata config tree in
          return ~tree ()
      | Profile collapse ->
          let* tree, inboxes, level =
            profile ~collapse level inboxes config extra tree
          in
          Lwt.return_ok (tree, inboxes, level)
      | Unknown s ->
          let*! () = Lwt_io.eprintf "Unknown command `%s`\n%!" s in
          return ()
      | Stop -> return ()
    in
    go command
end