type t = {
    localpart : string option;
    domainpart : string;
    resourcepart : string option;
  } [@@deriving show { with_path = false }]

exception InvalidUTF8

type uchars = Uchar.t list

let uchars_of_string s : uchars =
  let open Uchar in
  let len = String.length s in
  let rec loop acc i =
    if i >= len then
      List.rev acc
    else
      let c = String.get_utf_8_uchar s i in
      if not (utf_decode_is_valid c)
      then raise InvalidUTF8
      else
        let k = utf_decode_length c in
        let u = utf_decode_uchar c in
        loop (u :: acc) (i + k)
  in loop [] 0

let string_of_uchars (u : uchars) =
  let buf = Buffer.create (4 * List.length u) in
  List.iter (Buffer.add_utf_8_uchar buf) u;
  Buffer.to_bytes buf |> String.of_bytes

let of_string (s : string) : t =
  let open List in
  let jid = uchars_of_string s in
  let rest, resourcepart =
    find_mapi (fun i c ->
        if c = Uchar.of_char '/'
        then Some (take i jid, drop (i+1) jid |> string_of_uchars)
        else None) jid
    |> Option.fold ~none:(jid, None) ~some:(fun (rest, res) -> rest, Some res)
  in
  let localpart, domainpart =
    rev rest
    |> find_mapi (fun i c ->
           if c = Uchar.of_char '@'
           then let idx = length rest - i
                in Some (take (idx-1) rest |> string_of_uchars,
                         drop idx rest |> string_of_uchars)
           else None)
    |> Option.fold ~none:(None, rest |> string_of_uchars) ~some:(fun (loc, dom) -> Some loc, dom)
  in {
       localpart;
       resourcepart;
       domainpart;
     }

let%expect_test {|A "bare JID"|} =
  "juliet@example.com" |> of_string |> show |> print_string;
  [%expect {|
    { localpart = (Some "juliet"); domainpart = "example.com";
      resourcepart = None } |}]

let%expect_test {|A "full JID"|} =
  "juliet@example.com/foo" |> of_string |> show |> print_string;
  [%expect {|
    { localpart = (Some "juliet"); domainpart = "example.com";
      resourcepart = (Some "foo") } |}]

let to_string ({ localpart; domainpart; resourcepart } : t) =
  let local = Option.fold ~none:"" ~some:(fun l -> l ^ "@") localpart
  and resource = Option.fold ~none:"" ~some:(fun r -> "/" ^ r) resourcepart
  in local ^ domainpart ^ resource