(*---------------------------------------------------------------------------
  Copyright (c) 2025 Anil Madhavapeddy <anil@recoil.org>. 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 - represents how a token is interpreted in context *)
module Index = struct
  type t = [ `Mem of string | `Nth of int | `End ]

  let pp ppf = function
    | `Mem s -> Format.fprintf ppf "/%s" (Token.escape s)
    | `Nth n -> Format.fprintf ppf "/%d" n
    | `End -> Format.fprintf ppf "/-"

  let equal i1 i2 = match i1, i2 with
    | `Mem s1, `Mem s2 -> String.equal s1 s2
    | `Nth n1, `Nth n2 -> Int.equal n1 n2
    | `End, `End -> true
    | _ -> false

  let compare i1 i2 = match i1, i2 with
    | `Mem s1, `Mem s2 -> String.compare s1 s2
    | `Mem _, _ -> -1
    | _, `Mem _ -> 1
    | `Nth n1, `Nth n2 -> Int.compare n1 n2
    | `Nth _, `End -> -1
    | `End, `Nth _ -> 1
    | `End, `End -> 0

  let of_path_index (idx : Jsont.Path.index) : t =
    match idx with
    | Jsont.Path.Mem (s, _meta) -> `Mem s
    | Jsont.Path.Nth (n, _meta) -> `Nth n

  let to_path_index (idx : t) : Jsont.Path.index option =
    match idx with
    | `Mem s -> Some (Jsont.Path.Mem (s, Jsont.Meta.none))
    | `Nth n -> Some (Jsont.Path.Nth (n, Jsont.Meta.none))
    | `End -> None
end

(* Internal representation: raw unescaped tokens.
   Per RFC 6901, interpretation as member name vs array index
   depends on the JSON value type at evaluation time. *)
module Segment = struct
  type t =
    | Token of string  (* Unescaped reference token *)
    | End              (* The "-" token for end-of-array *)

  let of_escaped_string s =
    if s = "-" then End
    else Token (Token.unescape s)

  let to_escaped_string = function
    | Token s -> Token.escape s
    | End -> "-"

  (* Convert to Index for a given JSON value type *)
  let to_index seg ~for_array =
    match seg with
    | End -> `End
    | Token s ->
      if for_array then
        match Token.is_valid_array_index s with
        | Some n -> `Nth n
        | None -> `Mem s  (* Invalid index becomes member for error msg *)
      else
        `Mem s

  (* Convert from Index *)
  let of_index = function
    | `End -> End
    | `Mem s -> Token s
    | `Nth n -> Token (string_of_int n)
end

(* Pointer type - list of segments *)
type t = Segment.t list

let root = []

let is_root p = p = []

(* Convert indices to segments *)
let make indices = List.map Segment.of_index indices

(* Convert segments to indices, assuming array context for numeric tokens *)
let indices p = List.map (fun seg -> Segment.to_index seg ~for_array:true) p

let append p idx = p @ [Segment.of_index idx]

let concat p1 p2 = p1 @ p2

let parent p = match List.rev p with
  | [] -> None
  | _ :: rest -> Some (List.rev rest)

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

(* Parsing *)

let of_string s =
  if s = "" then root
  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_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 =
  of_string (percent_decode s)

let of_uri_fragment_result s =
  try Ok (of_uri_fragment s)
  with Jsont.Error e -> Error (Jsont.Error.to_string e)

(* Serialization *)

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_escaped_string seg)
    ) p;
    Buffer.contents b

(* URI fragment percent-encoding *)
let needs_percent_encoding c =
  (* RFC 3986 fragment: unreserved / pct-encoded / sub-delims / ":" / "@" / "/" / "?" *)
  (* unreserved = ALPHA / DIGIT / "-" / "." / "_" / "~" *)
  (* sub-delims = "!" / "$" / "&" / "'" / "(" / ")" / "*" / "+" / "," / ";" / "=" *)
  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 ppf p =
  let pp_index ppf = function
    | `Mem s -> Format.fprintf ppf "`Mem %S" s
    | `Nth n -> Format.fprintf ppf "`Nth %d" n
    | `End -> Format.fprintf ppf "`End"
  in
  Format.fprintf ppf "[%a]"
    (Format.pp_print_list ~pp_sep:(fun ppf () -> Format.fprintf ppf "; ") pp_index)
    (indices p)

(* Comparison *)

let segment_equal s1 s2 = match s1, s2 with
  | Segment.Token t1, Segment.Token t2 -> String.equal t1 t2
  | Segment.End, Segment.End -> true
  | _ -> false

let segment_compare s1 s2 = match s1, s2 with
  | Segment.Token t1, Segment.Token t2 -> String.compare t1 t2
  | Segment.Token _, Segment.End -> -1
  | Segment.End, Segment.Token _ -> 1
  | Segment.End, Segment.End -> 0

let equal p1 p2 =
  List.equal segment_equal p1 p2

let compare p1 p2 =
  List.compare segment_compare p1 p2

(* Path conversion *)

let segment_of_path_index (idx : Jsont.Path.index) : Segment.t =
  match idx with
  | Jsont.Path.Mem (s, _meta) -> Segment.Token s
  | Jsont.Path.Nth (n, _meta) -> Segment.Token (string_of_int n)

let of_path (p : Jsont.Path.t) : t =
  List.rev_map segment_of_path_index (Jsont.Path.rev_indices p)

let to_path p =
  let rec convert acc = function
    | [] -> Some acc
    | Segment.End :: _ -> None
    | Segment.Token s :: rest ->
      (* For path conversion, we need to decide if it's a member or index.
         We use array context for numeric tokens since Jsont.Path distinguishes. *)
      let acc' = match Token.is_valid_array_index s with
        | Some n -> Jsont.Path.nth ~meta:Jsont.Meta.none n acc
        | None -> Jsont.Path.mem ~meta:Jsont.Meta.none s acc
      in
      convert acc' rest
  in
  convert Jsont.Path.root p

let to_path_exn p =
  match to_path p with
  | Some path -> path
  | None ->
    Jsont.Error.msgf Jsont.Meta.none
      "Cannot convert JSON Pointer with '-' index to Jsont.Path"

(* 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 *)

let rec eval_get p json =
  match p with
  | [] -> json
  | Segment.End :: _ ->
    Jsont.Error.msgf (Jsont.Json.meta json)
      "JSON Pointer: '-' (end marker) refers to nonexistent array element"
  | Segment.Token token :: rest ->
    (match json with
     | Jsont.Object (members, _) ->
       (* For objects, token is always a member name *)
       (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, _) ->
       (* For arrays, token must be a valid array index *)
       (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 json = eval_get p 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 *)

let rec eval_set p value json =
  match p with
  | [] -> value
  | [Segment.End] ->
    (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))
  | Segment.End :: _ ->
    Jsont.Error.msgf (Jsont.Json.meta json)
      "JSON Pointer: '-' (end marker) refers to nonexistent array element"
  | [Segment.Token token] ->
    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))
  | Segment.Token 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 rest 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 rest value child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let set p json ~value = eval_set p value json

(* Mutation: add (RFC 6902 semantics) *)

let rec eval_add p value json =
  match p with
  | [] -> value
  | [Segment.End] ->
    (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))
  | Segment.End :: _ ->
    Jsont.Error.msgf (Jsont.Json.meta json)
      "JSON Pointer: '-' in non-final position"
  | [Segment.Token token] ->
    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))
  | Segment.Token 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 rest 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 rest value child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let add p json ~value = eval_add p value json

(* Mutation: remove *)

let rec eval_remove p json =
  match p with
  | [] ->
    Jsont.Error.msgf Jsont.Meta.none "JSON Pointer: cannot remove root document"
  | [Segment.End] ->
    Jsont.Error.msgf (Jsont.Json.meta json)
      "JSON Pointer: '-' refers to nonexistent element"
  | Segment.End :: _ ->
    Jsont.Error.msgf (Jsont.Json.meta json)
      "JSON Pointer: '-' in non-final position"
  | [Segment.Token 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))
  | Segment.Token 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 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 rest child) elements, meta)
        | None -> error_index_out_of_bounds json n)
      ~on_other:(fun () -> error_cannot_navigate json)

let remove p json = eval_remove p json

(* Mutation: replace *)

let replace p json ~value =
  (* Replace requires the target to exist, unlike add *)
  let _ = get p json in (* Will raise if not found *)
  eval_set p value json

(* Mutation: move *)

let rec is_prefix_of p1 p2 =
  match p1, p2 with
  | [], _ -> true
  | _, [] -> false
  | h1 :: t1, h2 :: t2 -> segment_equal h1 h2 && is_prefix_of t1 t2

let move ~from ~path json =
  (* Check for cycle: path cannot be a proper prefix of from *)
  if is_prefix_of path from && not (equal path from) 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 ~path json =
  let value = get from json in
  add path json ~value

(* Mutation: test *)

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

(* Jsont codec *)

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:"JSON Pointer"
    ~doc:"RFC 6901 JSON Pointer"
    ~dec ~enc ())

let jsont_uri_fragment : t Jsont.t =
  let dec _meta s = of_uri_fragment s in
  let enc 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 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 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 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 p json ~value:re_encoded
  in
  Jsont.map Jsont.json ~dec ~enc:(fun j -> j)

let delete_path ?(allow_absent = false) p =
  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)