(*---------------------------------------------------------------------------
   Copyright (c) 2025 Anil Madhavapeddy. All rights reserved.
   SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

type result_reference = {
  result_of : string;
  name : string;
  path : Json_pointer.Jmap.t;
}

let result_reference ~result_of ~name ~path =
  { result_of; name; path }

let result_reference_of_strings ~result_of ~name ~path =
  { result_of; name; path = Json_pointer.Jmap.of_string path }

let result_reference_make result_of name path =
  { result_of; name; path }

let result_reference_jsont =
  let kind = "ResultReference" in
  Jsont.Object.map ~kind result_reference_make
  |> Jsont.Object.mem "resultOf" Jsont.string ~enc:(fun r -> r.result_of)
  |> Jsont.Object.mem "name" Jsont.string ~enc:(fun r -> r.name)
  |> Jsont.Object.mem "path" Json_pointer.Jmap.jsont ~enc:(fun r -> r.path)
  |> Jsont.Object.finish

(* Result reference resolution *)

let find_response ref ~responses =
  List.find_map (fun (call_id, name, json) ->
    if call_id = ref.result_of && name = ref.name
    then Some json
    else None
  ) responses

let resolve ref ~responses codec =
  match find_response ref ~responses with
  | None ->
    Jsont.Error.msgf Jsont.Meta.none
      "Result reference: no response found for resultOf=%s name=%s"
      ref.result_of ref.name
  | Some json ->
    let extraction_codec = Json_pointer.Jmap.path ref.path codec in
    match Jsont.Json.decode' extraction_codec json with
    | Ok v -> v
    | Error e -> raise (Jsont.Error e)

let resolve_ids ref ~responses =
  resolve ref ~responses (Jsont.list Jsont.string)

type t = {
  name : string;
  arguments : Jsont.json;
  method_call_id : string;
}

let create ~name ~arguments ~method_call_id =
  { name; arguments; method_call_id }

let name t = t.name
let arguments t = t.arguments
let method_call_id t = t.method_call_id

(* Helper to encode a typed value back to Jsont.json *)
let encode_json_value jsont value =
  match Jsont.Json.encode' jsont value with
  | Ok json -> json
  | Error _ -> Jsont.Object ([], Jsont.Meta.none)

let jsont =
  let kind = "Invocation" in
  (* Invocation is [name, args, callId] - a 3-element heterogeneous array *)
  (* We need to handle this as a json array since elements have different types *)
  let dec json =
    match json with
    | Jsont.Array ([name_json; arguments; call_id_json], _) ->
      let name = match name_json with
        | Jsont.String (s, _) -> s
        | _ -> Jsont.Error.msg Jsont.Meta.none "Invocation[0] must be a string"
      in
      let method_call_id = match call_id_json with
        | Jsont.String (s, _) -> s
        | _ -> Jsont.Error.msg Jsont.Meta.none "Invocation[2] must be a string"
      in
      { name; arguments; method_call_id }
    | Jsont.Array _ ->
      Jsont.Error.msg Jsont.Meta.none "Invocation must be a 3-element array"
    | _ ->
      Jsont.Error.msg Jsont.Meta.none "Invocation must be an array"
  in
  let enc t =
    Jsont.Array ([
      Jsont.String (t.name, Jsont.Meta.none);
      t.arguments;
      Jsont.String (t.method_call_id, Jsont.Meta.none);
    ], Jsont.Meta.none)
  in
  Jsont.map ~kind ~dec ~enc Jsont.json

let make_get ~method_call_id ~method_name args =
  let arguments = encode_json_value Proto_method.get_args_jsont args in
  { name = method_name; arguments; method_call_id }

let make_changes ~method_call_id ~method_name args =
  let arguments = encode_json_value Proto_method.changes_args_jsont args in
  { name = method_name; arguments; method_call_id }

let make_query (type f) ~method_call_id ~method_name
    ~(filter_cond_jsont : f Jsont.t) (args : f Proto_method.query_args) =
  let arguments = encode_json_value (Proto_method.query_args_jsont filter_cond_jsont) args in
  { name = method_name; arguments; method_call_id }