ocaml-json-pointer / src / json_pointer.ml

(*---------------------------------------------------------------------------
  Copyright (c) 2025 Anil Madhavapeddy <[e-mail]>. All rights reserved.
  SPDX-License-Identifier: ISC
 ---------------------------------------------------------------------------*)

(* Token escaping/unescaping per RFC 6901 Section 3-4 *)
module Token = struct
  type t = string

  let escape s =
    let b = Buffer.create (String.length s) in
    String.iter (function
      | '~' -> Buffer.add_string b "~0"
      | '/' -> Buffer.add_string b "~1"
      | c -> Buffer.add_char b c
    ) s;
    Buffer.contents b

  let unescape s =
    let len = String.length s in
    let b = Buffer.create len in
    let rec loop i =
      if i >= len then Buffer.contents b
      else match s.[i] with
        | '~' when i + 1 >= len ->
          Jsont.Error.msgf Jsont.Meta.none
            "Invalid JSON Pointer: incomplete escape sequence at end"
        | '~' ->
          (match s.[i + 1] with
           | '0' -> Buffer.add_char b '~'; loop (i + 2)
           | '1' -> Buffer.add_char b '/'; loop (i + 2)
           | c ->
             Jsont.Error.msgf Jsont.Meta.none
               "Invalid JSON Pointer: invalid escape sequence ~%c" c)
        | c -> Buffer.add_char b c; loop (i + 1)
    in
    loop 0

  (* Check if a token is a valid array index per RFC 6901 ABNF:
     array-index = %x30 / ( %x31-39 *(%x30-39) )
     i.e., "0" or a non-zero digit followed by any digits *)
  let is_valid_array_index s =
    let len = String.length s in
    let is_digit c = c >= '0' && c <= '9' in
    if len = 0 then None
    else if len = 1 && s.[0] = '0' then Some 0
    else if s.[0] >= '1' && s.[0] <= '9' then
      let rec all_digits i =
        if i >= len then true
        else if is_digit s.[i] then all_digits (i + 1)
        else false
      in
      if all_digits 1 then int_of_string_opt s else None
    else None
end

(* Index type - directly reuses Jsont.Path.index *)
type index = Jsont.Path.index

(* Convenience constructors *)
let mem ?(meta = Jsont.Meta.none) s : index = Jsont.Path.Mem (s, meta)
let nth ?(meta = Jsont.Meta.none) n : index = Jsont.Path.Nth (n, meta)

let pp_index ppf = function
  | Jsont.Path.Mem (s, _) -> Format.fprintf ppf "/%s" (Token.escape s)
  | Jsont.Path.Nth (n, _) -> Format.fprintf ppf "/%d" n

let equal_index i1 i2 = match i1, i2 with
  | Jsont.Path.Mem (s1, _), Jsont.Path.Mem (s2, _) -> String.equal s1 s2
  | Jsont.Path.Nth (n1, _), Jsont.Path.Nth (n2, _) -> Int.equal n1 n2
  | _ -> false

let compare_index i1 i2 = match i1, i2 with
  | Jsont.Path.Mem (s1, _), Jsont.Path.Mem (s2, _) -> String.compare s1 s2
  | Jsont.Path.Mem _, Jsont.Path.Nth _ -> -1
  | Jsont.Path.Nth _, Jsont.Path.Mem _ -> 1
  | Jsont.Path.Nth (n1, _), Jsont.Path.Nth (n2, _) -> Int.compare n1 n2

(* Internal representation: raw unescaped tokens *)
module Segment = struct
  type t = string  (* Unescaped reference token *)

  let of_escaped_string s = Token.unescape s

  let to_escaped_string s = Token.escape s

  let of_index = function
    | Jsont.Path.Mem (s, _) -> s
    | Jsont.Path.Nth (n, _) -> string_of_int n

  let to_index s : index =
    match Token.is_valid_array_index s with
    | Some n -> nth n
    | None -> mem s
end

(* Phantom types *)
type nav
type append

(* Pointer type with phantom type parameter *)
type _ t = {
  segments : Segment.t list;
  is_append : bool;  (* true if ends with "-" *)
}

(* Existential wrapper *)
type any = Any : _ t -> any

let root = { segments = []; is_append = false }

let is_root p = p.segments = [] && not p.is_append

let make indices =
  { segments = List.map Segment.of_index indices; is_append = false }

let ( / ) p idx =
  { segments = p.segments @ [Segment.of_index idx]; is_append = false }

let append_index = ( / )

let at_end p =
  { segments = p.segments; is_append = true }

let concat p1 p2 =
  { segments = p1.segments @ p2.segments; is_append = false }

let parent p =
  match List.rev p.segments with
  | [] -> None
  | _ :: rest -> Some { segments = List.rev rest; is_append = false }

let last p =
  match List.rev p.segments with
  | [] -> None
  | seg :: _ -> Some (Segment.to_index seg)

let indices (type a) (p : a t) = List.map Segment.to_index p.segments

(* Coercion and inspection *)

let any (type a) (p : a t) : any = Any p

let is_nav (Any p) = not p.is_append

let to_nav (Any p) =
  if p.is_append then None
  else Some { segments = p.segments; is_append = false }

let to_nav_exn (Any p) =
  if p.is_append then
    Jsont.Error.msgf Jsont.Meta.none
      "JSON Pointer: cannot convert append pointer to nav pointer"
  else
    { segments = p.segments; is_append = false }

(* Parsing *)

let parse_segments s =
  if s = "" then []
  else if s.[0] <> '/' then
    Jsont.Error.msgf Jsont.Meta.none
      "Invalid JSON Pointer: must be empty or start with '/': %s" s
  else
    let rest = String.sub s 1 (String.length s - 1) in
    let tokens = String.split_on_char '/' rest in
    List.map Segment.of_escaped_string tokens

let of_string_kind s : [ `Nav of nav t | `Append of append t ] =
  let segments = parse_segments s in
  (* Check if ends with "-" *)
  match List.rev segments with
  | "-" :: rest ->
    (* Validate that "-" only appears at the end *)
    if List.exists (( = ) "-") rest then
      Jsont.Error.msgf Jsont.Meta.none
        "Invalid JSON Pointer: '-' can only appear at the end";
    `Append { segments = List.rev rest; is_append = true }
  | _ ->
    (* Validate no "-" anywhere *)
    if List.exists (( = ) "-") segments then
      Jsont.Error.msgf Jsont.Meta.none
        "Invalid JSON Pointer: '-' can only appear at the end";
    `Nav { segments; is_append = false }

let of_string s : any =
  match of_string_kind s with
  | `Nav p -> Any p
  | `Append p -> Any p

let of_string_nav s : nav t =
  match of_string_kind s with
  | `Nav p -> p
  | `Append _ ->
    Jsont.Error.msgf Jsont.Meta.none
      "Invalid JSON Pointer: '-' not allowed in navigation pointer"

let of_string_result s =
  try Ok (of_string s)
  with Jsont.Error e -> Error (Jsont.Error.to_string e)

(* URI fragment percent-decoding *)
let hex_value c =
  if c >= '0' && c <= '9' then Char.code c - Char.code '0'
  else if c >= 'A' && c <= 'F' then Char.code c - Char.code 'A' + 10
  else if c >= 'a' && c <= 'f' then Char.code c - Char.code 'a' + 10
  else -1

let percent_decode s =
  let len = String.length s in
  let b = Buffer.create len in
  let rec loop i =
    if i >= len then Buffer.contents b
    else match s.[i] with
      | '%' when i + 2 < len ->
        let h1 = hex_value s.[i + 1] in
        let h2 = hex_value s.[i + 2] in
        if h1 >= 0 && h2 >= 0 then begin
          Buffer.add_char b (Char.chr ((h1 lsl 4) lor h2));
          loop (i + 3)
        end else
          Jsont.Error.msgf Jsont.Meta.none
            "Invalid percent-encoding at position %d" i
      | '%' ->
        Jsont.Error.msgf Jsont.Meta.none
          "Incomplete percent-encoding at position %d" i
      | c -> Buffer.add_char b c; loop (i + 1)
  in
  loop 0

let of_uri_fragment s : any = of_string (percent_decode s)

let of_uri_fragment_nav s = of_string_nav (percent_decode s)

let of_uri_fragment_result s : (any, string) result =
  try Ok (of_uri_fragment s)
  with Jsont.Error e -> Error (Jsont.Error.to_string e)

(* Serialization *)

let to_string (type a) (p : a t) =
  let base =
    if p.segments = [] then ""
    else
      let b = Buffer.create 64 in
      List.iter (fun seg ->
        Buffer.add_char b '/';
        Buffer.add_string b (Segment.to_escaped_string seg)
      ) p.segments;
      Buffer.contents b
  in
  if p.is_append then base ^ "/-" else base

(* URI fragment percent-encoding *)
let needs_percent_encoding c =
  not (
    (c >= 'A' && c <= 'Z') ||
    (c >= 'a' && c <= 'z') ||
    (c >= '0' && c <= '9') ||
    c = '-' || c = '.' || c = '_' || c = '~' ||
    c = '!' || c = '$' || c = '&' || c = '\'' ||
    c = '(' || c = ')' || c = '*' || c = '+' ||
    c = ',' || c = ';' || c = '=' ||
    c = ':' || c = '@' || c = '/' || c = '?'
  )

let hex_char n =
  if n < 10 then Char.chr (Char.code '0' + n)
  else Char.chr (Char.code 'A' + n - 10)

let percent_encode s =
  let b = Buffer.create (String.length s * 3) in
  String.iter (fun c ->
    if needs_percent_encoding c then begin
      let code = Char.code c in
      Buffer.add_char b '%';
      Buffer.add_char b (hex_char (code lsr 4));
      Buffer.add_char b (hex_char (code land 0xF))
    end else
      Buffer.add_char b c
  ) s;
  Buffer.contents b

let to_uri_fragment p = percent_encode (to_string p)

let pp ppf p = Format.pp_print_string ppf (to_string p)

let pp_verbose (type a) ppf (p : a t) =
  let pp_idx ppf seg =
    match Token.is_valid_array_index seg with
    | Some n -> Format.fprintf ppf "Nth %d" n
    | None -> Format.fprintf ppf {|Mem "%s"|} seg
  in
  Format.fprintf ppf "[%a]%s"
    (Format.pp_print_list ~pp_sep:(fun ppf () -> Format.fprintf ppf "; ") pp_idx)
    p.segments
    (if p.is_append then " /-" else "")

(* Comparison *)

let equal (type a b) (p1 : a t) (p2 : b t) =
  List.equal String.equal p1.segments p2.segments &&
  p1.is_append = p2.is_append

let compare (type a b) (p1 : a t) (p2 : b t) =
  match List.compare String.compare p1.segments p2.segments with
  | 0 -> Bool.compare p1.is_append p2.is_append
  | n -> n

(* Path conversion *)

let of_path (p : Jsont.Path.t) : nav t =
  let segments = List.rev_map Segment.of_index (Jsont.Path.rev_indices p) in
  { segments; is_append = false }

let to_path (p : nav t) : Jsont.Path.t =
  List.fold_left (fun acc seg ->
    match Token.is_valid_array_index seg with
    | Some n -> Jsont.Path.nth n acc
    | None -> Jsont.Path.mem seg acc
  ) Jsont.Path.root p.segments

(* Evaluation helpers *)

let json_sort_string (j : Jsont.json) =
  match j with
  | Null _ -> "null"
  | Bool _ -> "boolean"
  | Number _ -> "number"
  | String _ -> "string"
  | Array _ -> "array"
  | Object _ -> "object"

let get_member name (obj : Jsont.object') =
  List.find_opt (fun ((n, _), _) -> String.equal n name) obj

let get_nth n (arr : Jsont.json list) =
  if n < 0 || n >= List.length arr then None
  else Some (List.nth arr n)

(* Evaluation - only for nav pointers *)

let rec eval_get segments json =
  match segments with
  | [] -> json
  | token :: rest ->
    (match json with
     | Jsont.Object (members, _) ->
       (match get_member token members with
        | Some (_, value) -> eval_get rest value
        | None ->
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: member '%s' not found" token)
     | Jsont.Array (elements, _) ->
       (match Token.is_valid_array_index token with
        | Some n ->
          (match get_nth n elements with
           | Some value -> eval_get rest value
           | None ->
             Jsont.Error.msgf (Jsont.Json.meta json)
               "JSON Pointer: index %d out of bounds (array has %d elements)"
               n (List.length elements))
        | None ->
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: invalid array index '%s'" token)
     | _ ->
       Jsont.Error.msgf (Jsont.Json.meta json)
         "JSON Pointer: cannot index into %s with '%s'"
         (json_sort_string json) token)

let get (p : nav t) json = eval_get p.segments json

let get_result p json =
  try Ok (get p json)
  with Jsont.Error e -> Error e

let find p json =
  try Some (get p json)
  with Jsont.Error _ -> None

(* Mutation helpers *)

let set_member name value (obj : Jsont.object') : Jsont.object' =
  let rec loop found acc = function
    | [] ->
      if found then List.rev acc
      else List.rev_append acc [((name, Jsont.Meta.none), value)]
    | ((n, m), _) :: rest when String.equal n name ->
      loop true (((n, m), value) :: acc) rest
    | mem :: rest ->
      loop found (mem :: acc) rest
  in
  loop false [] obj

let remove_member name (obj : Jsont.object') : Jsont.object' =
  List.filter (fun ((n, _), _) -> not (String.equal n name)) obj

let insert_at n value lst =
  let rec loop i acc = function
    | rest when i = n -> List.rev_append acc (value :: rest)
    | [] -> List.rev acc
    | h :: t -> loop (i + 1) (h :: acc) t
  in
  loop 0 [] lst

let remove_at n lst =
  List.filteri (fun i _ -> i <> n) lst

let replace_at n value lst =
  List.mapi (fun i v -> if i = n then value else v) lst

(* Common navigation for mutation operations *)

let navigate_to_child token json ~on_object ~on_array ~on_other =
  match json with
  | Jsont.Object (members, meta) -> on_object members meta
  | Jsont.Array (elements, meta) ->
    (match Token.is_valid_array_index token with
     | Some n -> on_array elements meta n
     | None ->
       Jsont.Error.msgf (Jsont.Json.meta json)
         "JSON Pointer: invalid array index '%s'" token)
  | _ -> on_other ()

let error_member_not_found json token =
  Jsont.Error.msgf (Jsont.Json.meta json) "JSON Pointer: member '%s' not found" token

let error_index_out_of_bounds json n =
  Jsont.Error.msgf (Jsont.Json.meta json) "JSON Pointer: index %d out of bounds" n

let error_cannot_navigate json =
  Jsont.Error.msgf (Jsont.Json.meta json)
    "JSON Pointer: cannot navigate through %s" (json_sort_string json)

(* Mutation: set - works with any pointer type *)

let rec eval_set_segments segments is_append value json =
  match segments, is_append with
  | [], false -> value
  | [], true ->
    (* Append to array *)
    (match json with
     | Jsont.Array (elements, meta) -> Jsont.Array (elements @ [value], meta)
     | _ ->
       Jsont.Error.msgf (Jsont.Json.meta json)
         "JSON Pointer: '-' can only be used on arrays, got %s"
         (json_sort_string json))
  | [token], false ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        if Option.is_some (get_member token members) then
          Jsont.Object (set_member token value members, meta)
        else
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: member '%s' not found for set" token)
      ~on_array:(fun elements meta n ->
        if n < List.length elements then
          Jsont.Array (replace_at n value elements, meta)
        else
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: index %d out of bounds for set" n)
      ~on_other:(fun () ->
        Jsont.Error.msgf (Jsont.Json.meta json)
          "JSON Pointer: cannot set in %s" (json_sort_string json))
  | [token], true ->
    (* Navigate to token, then append *)
    navigate_to_child token json
      ~on_object:(fun members meta ->
        match get_member token members with
        | Some (_, child) ->
          let child' = eval_set_segments [] true value child in
          Jsont.Object (set_member token child' members, meta)
        | None -> error_member_not_found json token)
      ~on_array:(fun elements meta n ->
        match get_nth n elements with
        | Some child ->
          let child' = eval_set_segments [] true value child in
          Jsont.Array (replace_at n child' elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)
  | token :: rest, _ ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        match get_member token members with
        | Some (_, child) ->
          Jsont.Object (set_member token (eval_set_segments rest is_append value child) members, meta)
        | None -> error_member_not_found json token)
      ~on_array:(fun elements meta n ->
        match get_nth n elements with
        | Some child ->
          Jsont.Array (replace_at n (eval_set_segments rest is_append value child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let set (Any p) json ~value =
  eval_set_segments p.segments p.is_append value json

(* Mutation: add (RFC 6902 semantics) - works with any pointer type *)

let rec eval_add_segments segments is_append value json =
  match segments, is_append with
  | [], false -> value
  | [], true ->
    (* Append to array *)
    (match json with
     | Jsont.Array (elements, meta) -> Jsont.Array (elements @ [value], meta)
     | _ ->
       Jsont.Error.msgf (Jsont.Json.meta json)
         "JSON Pointer: '-' can only be used on arrays, got %s"
         (json_sort_string json))
  | [token], false ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        Jsont.Object (set_member token value members, meta))
      ~on_array:(fun elements meta n ->
        let len = List.length elements in
        if n <= len then
          Jsont.Array (insert_at n value elements, meta)
        else
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: index %d out of bounds for add (array has %d elements)"
            n len)
      ~on_other:(fun () ->
        Jsont.Error.msgf (Jsont.Json.meta json)
          "JSON Pointer: cannot add to %s" (json_sort_string json))
  | [token], true ->
    (* Navigate to token, then append *)
    navigate_to_child token json
      ~on_object:(fun members meta ->
        match get_member token members with
        | Some (_, child) ->
          let child' = eval_add_segments [] true value child in
          Jsont.Object (set_member token child' members, meta)
        | None -> error_member_not_found json token)
      ~on_array:(fun elements meta n ->
        match get_nth n elements with
        | Some child ->
          let child' = eval_add_segments [] true value child in
          Jsont.Array (replace_at n child' elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)
  | token :: rest, _ ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        match get_member token members with
        | Some (_, child) ->
          Jsont.Object (set_member token (eval_add_segments rest is_append value child) members, meta)
        | None -> error_member_not_found json token)
      ~on_array:(fun elements meta n ->
        match get_nth n elements with
        | Some child ->
          Jsont.Array (replace_at n (eval_add_segments rest is_append value child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let add (Any p) json ~value =
  eval_add_segments p.segments p.is_append value json

(* Mutation: remove - only for nav pointers *)

let rec eval_remove_segments segments json =
  match segments with
  | [] ->
    Jsont.Error.msgf Jsont.Meta.none "JSON Pointer: cannot remove root document"
  | [token] ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        if Option.is_some (get_member token members) then
          Jsont.Object (remove_member token members, meta)
        else
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: member '%s' not found for remove" token)
      ~on_array:(fun elements meta n ->
        if n < List.length elements then
          Jsont.Array (remove_at n elements, meta)
        else
          Jsont.Error.msgf (Jsont.Json.meta json)
            "JSON Pointer: index %d out of bounds for remove" n)
      ~on_other:(fun () ->
        Jsont.Error.msgf (Jsont.Json.meta json)
          "JSON Pointer: cannot remove from %s" (json_sort_string json))
  | token :: rest ->
    navigate_to_child token json
      ~on_object:(fun members meta ->
        match get_member token members with
        | Some (_, child) ->
          Jsont.Object (set_member token (eval_remove_segments rest child) members, meta)
        | None -> error_member_not_found json token)
      ~on_array:(fun elements meta n ->
        match get_nth n elements with
        | Some child ->
          Jsont.Array (replace_at n (eval_remove_segments rest child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let remove (p : nav t) json = eval_remove_segments p.segments json

(* Mutation: replace - only for nav pointers *)

let replace (p : nav t) json ~value =
  let _ = get p json in (* Will raise if not found *)
  eval_set_segments p.segments false value json

(* Mutation: move *)

let move ~(from : nav t) ~(path : any) json =
  let (Any p) = path in
  (* Check for cycle: path cannot be a proper prefix of from *)
  let from_segs = from.segments in
  let path_segs = p.segments in
  let rec is_prefix p1 p2 = match p1, p2 with
    | [], _ -> true
    | _, [] -> false
    | h1 :: t1, h2 :: t2 -> String.equal h1 h2 && is_prefix t1 t2
  in
  if is_prefix path_segs from_segs &&
     not (List.equal String.equal path_segs from_segs && p.is_append = false) then
    Jsont.Error.msgf Jsont.Meta.none
      "JSON Pointer: move would create cycle (path is prefix of from)";
  let value = get from json in
  let json' = remove from json in
  add path json' ~value

(* Mutation: copy *)

let copy ~(from : nav t) ~(path : any) json =
  let value = get from json in
  add path json ~value

(* Mutation: test *)

let test (p : nav t) json ~expected =
  Option.fold ~none:false ~some:(Jsont.Json.equal expected) (find p json)

(* Jsont codec *)

let jsont : any Jsont.t =
  let dec _meta s = of_string s in
  let enc (Any p) = to_string p in
  Jsont.Base.string (Jsont.Base.map
    ~kind:"JSON Pointer"
    ~doc:"RFC 6901 JSON Pointer"
    ~dec ~enc ())

let jsont_kind : [ `Nav of nav t | `Append of append t ] Jsont.t =
  let dec _meta s = of_string_kind s in
  let enc = function
    | `Nav p -> to_string p
    | `Append p -> to_string p
  in
  Jsont.Base.string (Jsont.Base.map
    ~kind:"JSON Pointer (kind)"
    ~doc:"RFC 6901 JSON Pointer with kind tag"
    ~dec ~enc ())

let jsont_nav : nav t Jsont.t =
  let dec _meta s = of_string_nav s in
  let enc p = to_string p in
  Jsont.Base.string (Jsont.Base.map
    ~kind:"JSON Pointer (nav)"
    ~doc:"RFC 6901 JSON Pointer (navigation only)"
    ~dec ~enc ())

let jsont_uri_fragment : any Jsont.t =
  let dec _meta s = of_uri_fragment s in
  let enc (Any p) = to_uri_fragment p in
  Jsont.Base.string (Jsont.Base.map
    ~kind:"JSON Pointer (URI fragment)"
    ~doc:"RFC 6901 JSON Pointer in URI fragment encoding"
    ~dec ~enc ())

(* Query combinators *)

let path ?absent (p : nav t) t =
  let dec json =
    match find p json with
    | Some value ->
      (match Jsont.Json.decode' t value with
       | Ok v -> v
       | Error e -> raise (Jsont.Error e))
    | None ->
      match absent with
      | Some v -> v
      | None ->
        Jsont.Error.msgf Jsont.Meta.none
          "JSON Pointer %s: path not found" (to_string p)
  in
  Jsont.map Jsont.json ~dec ~enc:(fun _ ->
    Jsont.Error.msgf Jsont.Meta.none "path: encode not supported")

let set_path ?(allow_absent = false) t (p : any) v =
  let encoded = match Jsont.Json.encode' t v with
    | Ok json -> json
    | Error e -> raise (Jsont.Error e)
  in
  let dec json =
    if allow_absent then
      add p json ~value:encoded
    else
      set p json ~value:encoded
  in
  Jsont.map Jsont.json ~dec ~enc:(fun j -> j)

let update_path ?absent (p : nav t) t =
  let dec json =
    let value = match find p json with
      | Some v -> v
      | None ->
        match absent with
        | Some v ->
          (match Jsont.Json.encode' t v with
           | Ok j -> j
           | Error e -> raise (Jsont.Error e))
        | None ->
          Jsont.Error.msgf Jsont.Meta.none
            "JSON Pointer %s: path not found" (to_string p)
    in
    let decoded = match Jsont.Json.decode' t value with
      | Ok v -> v
      | Error e -> raise (Jsont.Error e)
    in
    let re_encoded = match Jsont.Json.encode' t decoded with
      | Ok j -> j
      | Error e -> raise (Jsont.Error e)
    in
    set (Any p) json ~value:re_encoded
  in
  Jsont.map Jsont.json ~dec ~enc:(fun j -> j)

let delete_path ?(allow_absent = false) (p : nav t) =
  let dec json =
    if allow_absent then
      match find p json with
      | Some _ -> remove p json
      | None -> json
    else
      remove p json
  in
  Jsont.map Jsont.json ~dec ~enc:(fun j -> j)

(* JMAP Extended Pointers - RFC 8620 Section 3.7 *)
module Jmap = struct
  (* Extended segment type: regular tokens or wildcard *)
  type segment =
    | Token of string  (* Unescaped reference token *)
    | Wildcard         (* The * token for array mapping *)

  type t = segment list

  let parse_segments s =
    if s = "" then []
    else if s.[0] <> '/' then
      Jsont.Error.msgf Jsont.Meta.none
        "Invalid JMAP Pointer: must be empty or start with '/': %s" s
    else
      let rest = String.sub s 1 (String.length s - 1) in
      let tokens = String.split_on_char '/' rest in
      List.map (fun tok ->
        if tok = "*" then Wildcard
        else if tok = "-" then
          Jsont.Error.msgf Jsont.Meta.none
            "Invalid JMAP Pointer: '-' not supported in result reference paths"
        else Token (Token.unescape tok)
      ) tokens

  let of_string s = parse_segments s

  let of_string_result s =
    try Ok (of_string s)
    with Jsont.Error e -> Error (Jsont.Error.to_string e)

  let segment_to_string = function
    | Token s -> Token.escape s
    | Wildcard -> "*"

  let to_string p =
    if p = [] then ""
    else
      let b = Buffer.create 64 in
      List.iter (fun seg ->
        Buffer.add_char b '/';
        Buffer.add_string b (segment_to_string seg)
      ) p;
      Buffer.contents b

  let pp ppf p = Format.pp_print_string ppf (to_string p)

  (* Evaluation with wildcard support *)
  let rec eval_segments segments json =
    match segments with
    | [] -> json
    | Wildcard :: rest ->
      (* Wildcard: map through array, flatten results *)
      (match json with
       | Jsont.Array (elements, meta) ->
         let results = List.map (eval_segments rest) elements in
         (* Flatten: if a result is an array, inline its contents *)
         let flattened = List.concat_map (function
           | Jsont.Array (elems, _) -> elems
           | other -> [other]
         ) results in
         Jsont.Array (flattened, meta)
       | _ ->
         Jsont.Error.msgf (Jsont.Json.meta json)
           "JMAP Pointer: '*' can only be used on arrays, got %s"
           (json_sort_string json))
    | Token token :: rest ->
      (* Standard token: navigate into object or array *)
      (match json with
       | Jsont.Object (members, _) ->
         (match get_member token members with
          | Some (_, value) -> eval_segments rest value
          | None ->
            Jsont.Error.msgf (Jsont.Json.meta json)
              "JMAP Pointer: member '%s' not found" token)
       | Jsont.Array (elements, _) ->
         (match Token.is_valid_array_index token with
          | Some n ->
            (match get_nth n elements with
             | Some value -> eval_segments rest value
             | None ->
               Jsont.Error.msgf (Jsont.Json.meta json)
                 "JMAP Pointer: index %d out of bounds (array has %d elements)"
                 n (List.length elements))
          | None ->
            Jsont.Error.msgf (Jsont.Json.meta json)
              "JMAP Pointer: invalid array index '%s'" token)
       | _ ->
         Jsont.Error.msgf (Jsont.Json.meta json)
           "JMAP Pointer: cannot index into %s with '%s'"
           (json_sort_string json) token)

  let eval p json = eval_segments p json

  let eval_result p json =
    try Ok (eval p json)
    with Jsont.Error e -> Error e

  let find p json =
    try Some (eval p json)
    with Jsont.Error _ -> None

  let jsont : t Jsont.t =
    let dec _meta s = of_string s in
    let enc p = to_string p in
    Jsont.Base.string (Jsont.Base.map
      ~kind:"JMAP Pointer"
      ~doc:"RFC 8620 JMAP extended JSON Pointer"
      ~dec ~enc ())

  (* Query combinators *)

  let path ?absent p codec =
    let dec json =
      match find p json with
      | Some extracted ->
        (match Jsont.Json.decode' codec extracted with
         | Ok v -> v
         | Error e -> raise (Jsont.Error e))
      | None ->
        match absent with
        | Some v -> v
        | None ->
          Jsont.Error.msgf Jsont.Meta.none
            "JMAP Pointer %s: path not found" (to_string p)
    in
    Jsont.map Jsont.json ~dec ~enc:(fun _ ->
      Jsont.Error.msgf Jsont.Meta.none "Jmap.path: encode not supported")

  let path_list p elem_codec =
    path p (Jsont.list elem_codec)
end