open Lwt.Syntax
open Markup

exception ClosedStream
exception InvalidStanza of string

type auth_mechanism =
  | PLAIN
  | Unknown of string [@@deriving show { with_path = false }]

let parse_auth_mechanism = function
  | "PLAIN" -> PLAIN
  | other -> Unknown other

type stream_features = {
    sasl_mechanisms : auth_mechanism list;
    starttls : [`Required | `Optional | `None];
    unknown_features : Xml.element list;
  }

(** [next stream] is a promise containing a full stanza of the fragments of
    [stream]. *)
let next (stream : (signal, async) stream) : (signal, sync) stream Lwt.t =
  let traverse_stanza depth fragment =
    let depth = match fragment with
      | `Start_element _ -> depth + 1
      | `End_element -> depth - 1
      | _ -> depth
    in ([fragment], if depth = 0 then None else Some depth)
  in transform traverse_stanza 0 stream |> Markup_lwt.load

(** [get stream] is a promise containing a single Xml element of [stream]. *)
let get (stream : (signal, async) stream) : Xml.element Lwt.t =
  let* signal = next stream
  in match Xml.tree signal with
     | Some xml -> Lwt.return xml
     | None -> Lwt.fail (InvalidStanza (signal |> write_xml |> to_string))

(** [start domain] is a promise containing a Portal (stream * push) connected to the
    XMPP server [domain].

    Currently, it doesn't handle anything except the initial [<open/>] stanza. *)
let start domain : Portal.t Lwt.t =
  let* stream, _push = Portal.connect domain
  in let push = function
       | None -> Portal.close (stream, _push);
       | anything -> _push anything
     in let+ _id = Portal.header domain (stream, _push)
        in stream, push

(** [parse_features el] is a [stream_features] record with all the features of the
    [<stream:features>] stanza contained in [el]. *)
let parse_features (el : Xml.element) : stream_features =
  let open Xml in
  let open Either in
  let parse_mechanism_stanza = function
    | Left {local_name = "mechanism"; children = [Right mechanism]; _} ->
       Some (parse_auth_mechanism mechanism)
    | _ -> None
  in
  let parse_feature (acc : stream_features) (feature : Xml.element) : stream_features =
    let parse_mechanisms ch =
      List.filter_map parse_mechanism_stanza ch
    and parse_starttls = function
      | [Left {local_name="required"; _}] -> `Required
      | [] -> `Optional
      | _ -> raise (InvalidStanza (element_to_string el))
    in match feature.local_name with
       | "mechanisms" -> {acc with sasl_mechanisms=parse_mechanisms feature.children}
       | "starttls" ->  {acc with starttls=parse_starttls feature.children}
       | _ -> {acc with unknown_features = feature :: acc.unknown_features}
  in List.fold_left
       parse_feature
       {sasl_mechanisms=[]; starttls=`None; unknown_features=[]}
       (List.filter_map find_left el.children)