type t = {
    localpart : string option;
    domainpart : string;
    resourcepart : string option;
  }

exception InvalidUTF8

type uchars = (Uchar.t * int) 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, i) :: acc) (i + k)
  in loop [] 0

let of_string (jid : string) : t =
  let open List in
  let open String in
  let u_jid = uchars_of_string jid in
  let u_barejid, domain_end =
    find_mapi (fun u_i (c, i) ->
        if c = Uchar.of_char '/'
        then Some (u_i, i)
        else None) u_jid
    |> Option.fold ~none:(u_jid, None) ~some:(fun (u_i, i) -> take u_i u_jid, Some i)
  in
  let domain_begin =
    rev u_barejid
    |> find_map (fun (c, i) ->
           if c = Uchar.of_char '@'
           then Some (i+1)
           else None)
  in
  let open Option in {
      localpart = map (fun i -> sub jid 0 (i - 1)) domain_begin;
      domainpart =
        begin
          let b = (value ~default:0 domain_begin) in
          sub jid b ((value ~default:(length jid) domain_end) - b)
        end;
      resourcepart = map (fun i -> sub jid (i + 1) (length jid - i - 1)) domain_end;
    }

let show ({ localpart; domainpart; resourcepart } : t) =
  let show_opt = Option.fold ~none:"None" ~some:(Format.sprintf {|Some "%s"|})
  in Printf.sprintf
{|{ localpart = %s;
  domainpart = "%s";
  resourcepart = %s }|}
       (show_opt localpart) domainpart (show_opt resourcepart)

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