(*---------------------------------------------------------------------------
   Copyright (c) 2024 The jsont programmers. All rights reserved.
   SPDX-License-Identifier: CC0-1.0
  ---------------------------------------------------------------------------*)

(** JSON-RPC codec https://www.jsonrpc.org/ *)

(* JSON-RPC version *)

type jsonrpc = [ `V2 ]

let jsonrpc_codec = Json.Codec.enum [ ("2.0", `V2) ]

(* JSON-RPC identifiers *)

type id = [ `String of string | `Number of float | `Null ]

let id_codec : id Json.Codec.t =
  let null = Json.Codec.null `Null in
  let string =
    let dec s = `String s in
    let enc = function `String s -> s | _ -> assert false in
    Json.Codec.map ~dec ~enc Json.Codec.string
  in
  let number =
    let dec n = `Number n in
    let enc = function `Number n -> n | _ -> assert false in
    Json.Codec.map ~dec ~enc Json.Codec.number
  in
  let enc = function
    | `Null -> null
    | `String _ -> string
    | `Number _ -> number
  in
  Json.Codec.any ~dec_null:null ~dec_string:string ~dec_number:number ~enc ()

(* JSON-RPC request object *)

type params = Json.t (* An array or object *)

let params_codec =
  let enc = function
    | Json.Value.Object _ | Json.Value.Array _ -> Json.Codec.Value.t
    | j ->
        let meta = Json.Meta.none in
        let fnd = Json.Sort.to_string (Json.Value.sort j) in
        Json.Error.fail_expected meta "object or array" ~fnd
  in
  let kind = "JSON-RPC params" in
  Json.Codec.any ~kind ~dec_array:Json.Codec.Value.t
    ~dec_object:Json.Codec.Value.t ~enc ()

type request = {
  jsonrpc : jsonrpc;
  method' : string;
  params : params option;
  id : id option;
}

let request jsonrpc method' params id = { jsonrpc; method'; params; id }

let request_codec : request Json.Codec.t =
  Json.Codec.Object.map request
  |> Json.Codec.Object.member "jsonrpc" jsonrpc_codec ~enc:(fun r -> r.jsonrpc)
  |> Json.Codec.Object.member "method" Json.Codec.string ~enc:(fun r ->
      r.method')
  |> Json.Codec.Object.opt_member "params" params_codec ~enc:(fun r -> r.params)
  |> Json.Codec.Object.opt_member "id" id_codec ~enc:(fun r -> r.id)
  |> Json.Codec.Object.seal

(* JSON-RPC error objects *)

type error = { code : int; message : string; data : Json.t option }

let error code message data = { code; message; data }

let error_codec =
  Json.Codec.Object.map error
  |> Json.Codec.Object.member "code" Json.Codec.int ~enc:(fun e -> e.code)
  |> Json.Codec.Object.member "message" Json.Codec.string ~enc:(fun e ->
      e.message)
  |> Json.Codec.Object.opt_member "data" Json.Codec.Value.t ~enc:(fun e ->
      e.data)
  |> Json.Codec.Object.seal

(* JSON-RPC response object *)

type response = { jsonrpc : jsonrpc; value : (Json.t, error) result; id : id }

let response jsonrpc result error id : response =
  let err_both () =
    Json.Error.msgf Json.Meta.none "Both %a and %a members are defined"
      Format.pp_print_string "result" Format.pp_print_string "error"
  in
  let err_none () =
    Json.Error.msgf Json.Meta.none "Missing either %a or %a member"
      Format.pp_print_string "result" Format.pp_print_string "error"
  in
  match (result, error) with
  | Some result, None -> { jsonrpc; value = Ok result; id }
  | None, Some error -> { jsonrpc; value = Error error; id }
  | Some _, Some _ -> err_both ()
  | None, None -> err_none ()

let response_result r = match r.value with Ok v -> Some v | Error _ -> None
let response_error r = match r.value with Ok _ -> None | Error e -> Some e

let response_codec : response Json.Codec.t =
  Json.Codec.Object.map response
  |> Json.Codec.Object.member "jsonrpc" jsonrpc_codec ~enc:(fun r -> r.jsonrpc)
  |> Json.Codec.Object.opt_member "result" Json.Codec.Value.t
       ~enc:response_result
  |> Json.Codec.Object.opt_member "error" error_codec ~enc:response_error
  |> Json.Codec.Object.member "id" id_codec ~enc:(fun r -> r.id)
  |> Json.Codec.Object.seal