open Lwt.Syntax
open Lwt.Infix
open Js_of_ocaml
open Markup
let jss = Js.string
let sjs = Js.to_string

(** Opaque domain name type. Currently a string. *)
type domain = string

(** [socket] is a framed Lwt stream (and its push function) communicating with a
    websocket.

    Valid XMPP WebSocket subprotocol frames must be sent to the stream, because it
    directly exposes the websocket under. *)
type socket = {
    domain : domain;
    stream : string Lwt_stream.t;
    push : string option -> unit;
  }

let domain_of_string (s : string) : domain = s
let domain_to_string (s : domain) : string = s

type t = {
    stream : (signal, async) stream;
    push : (signal, sync) stream option -> unit;
    mutable _socket : socket;
  }

let xmlns = "urn:ietf:params:xml:ns:xmpp-framing"

exception MalformedStanza of Markup.Error.t

(* sic. XEP-0156: "host-meta files MUST be fetched only over HTTPS". I don't make the
   rules. *)
let well_known_of (domain : string) = "https://" ^ domain ^ "/.well-known/host-meta"

(** [ws_endpoint domain] is a promise containing the XMPP websocket endpoint
    associated with [domain], by using the domain's Web-host Metadata.

    This function uses XMLHttpRequest, so while it should work fine in the browser, in
    environments that don't provide this constructor (Node.js), there should be some
    sort of polyfill.

    Lastly, if [domain] doesn't provide a well-formed Web-host Metadata file, the
    function throws an exception. *)
let ws_endpoint (domain : string) : string Lwt.t =
  let open Markup in
  (* This ugly function extracts the href element from a Link tag's attributes if it's a
     websocket. *)
  let link_websocket = function
    | ((_, "rel"), "urn:xmpp:alt-connections:websocket") :: ((_, "href"), href) :: _
      | ((_, "href"), href) :: ((_, "rel"), "urn:xmpp:alt-connections:websocket") :: _
      -> Some href
    | _ -> None
  in let parse_xrd xrd =
       string xrd
       |> parse_xml
       |> signals
       |> tree ~element:(fun (_, name) attributes children ->
              match name with
              | "Link" -> link_websocket attributes
              | "XRD" -> List.find_map (fun x -> x) children
              | _ -> None
            )
       |> Option.join
     in let* host_meta =
          Js_of_ocaml_lwt.XmlHttpRequest.perform_raw_url (well_known_of domain)
        in match parse_xrd host_meta.content with
           | Some x -> Lwt.return x
           | None -> Lwt.fail_with (domain ^ ": no WebSocket endpoint in Web-host Metadata.")

(** [connect url] is a promise to a [socket] communicating with [url] using the XMPP
    protocol.

    If the websocket is closed server-side, the stream closes. *)
let connect (domain : domain) : socket Lwt.t =
  let ws_stream (url : string) =
    let open Lwt_stream in
    let stream, message = create () in
    let open Lwt_condition in
    let is_open = create () in
    let (ws : WebSockets.webSocket Js.t) =
      new%js WebSockets.webSocket_withProtocol (jss url) (jss "xmpp")
    in let push = function
         | Some msg -> ws##send (jss msg)
         | None -> ws##close
       in ws##.onclose := Dom.handler (fun _ -> message None; Js._true);
          ws##.onmessage := Dom.handler (fun x -> Some (sjs x##.data) |> message; Js._false);
          ws##.onopen := Dom.handler (fun _ -> signal is_open (); Js._false);
          let+ () = wait is_open in {domain; stream; push}
  in ws_endpoint domain >>= ws_stream

(** [close] is a [<close/>] stanza. *)
let close = {|<close xmlns="|} ^ xmlns ^ {|" />|} |> string |> parse_xml |> signals

(** [portal_of_socket] is a Portal communicating with the socket [socket] via the
    Websocket subprotocol.


    This function is a complex wrapper that accepts Markup.ml signals and sends framed
    XML stanzas to the underlying socket, with exactly one stanza per frame, according
    to RFC 7935.

    In essence, it exposes an identical interface to the original XMPP streamed
    protocol.

    Here's an ASCII rendered flow of the data through the various streams.
        / -> push -> mu_stream -> to_frames -> ws_push -> \
    function user                                     websocket
        \ <--- stream <--- markup_lwt <--- ws_stream <--- / *)
let portal_of_socket (ws : socket) : t =
  let open Markup_lwt in
  (* When sending a malformed stanza (one that Markup.ml doesn't like), a
     MalformedStanza exception is raised. *)
  let report _ err =
    Lwt.fail (MalformedStanza err)
  in
  (* Consumes a stream of Markup.ml signals into a series of frames sent to the
     WebSocket. *)
  let fragment_stream, fragment_push = Lwt_stream.create () in
  let stream = ws.stream |> lwt_stream |> strings_to_bytes |> parse_xml ~report |> signals
  in
  let push = function
    | Some fs -> Markup.iter (fun f -> fragment_push (Some f)) fs
    | None -> begin
        (* We need to send the [<close>] stanza in full to the WebSocket. *)
        Markup.iter (fun f -> fragment_push (Some f)) close;
        Lwt.async
          (fun () ->
            (* We drain completely the stream when closing, so the socket can close. *)
            let+ () = Markup_lwt.drain stream
            in fragment_push None)
      end
  in
  (* Elements filters all elements based on the `Start_element. By simply making it
     return true every time, we get a stream of stream of elements, with each sub-stream
     being a full frame. *)
  Lwt.async
    (fun () -> fragment_stream
               |> lwt_stream
               |> elements (fun _ _ -> true)
               |> map (fun x -> write_xml x |> to_string)
               |> to_lwt_stream
               |> Lwt_stream.iter (fun s -> ws.push (Some s)));
  {stream; push; _socket=ws}

(** [stream socket] is a promise to a Portal connected to [domain].

    If [from] is specified, the <open /> stanza sent to the portal has a from
    parameter. *)
let stream ?from (socket : socket) : t Lwt.t =
  let portal = portal_of_socket socket in
  let {stream; push; _} = portal in
  let stanza =
    let attributes =
      let open Option in
      [(("", "xmlns"), xmlns);
       (("", "to"), socket.domain)]
      @ (map (fun jid -> (("", "from"), jid)) from |> to_list)
      @ [(("", "version"), "1.0")]
    in
    [`Start_element
       ((xmlns, "open"),
        attributes);
     `End_element]
  in push (Some (of_list stanza));
     let some_id ((_, name), value) = if name = "id" then Some value else None in
     let* stanza_open = Markup_lwt.next stream in
     let* id = match stanza_open with
       | Some `Start_element ((ns, "open"), attributes) when ns = xmlns ->
          let* close = Markup_lwt.next stream in
          begin match close with
          | Some `End_element -> List.find_map some_id attributes |> Lwt.return
          | _ -> Lwt.return_none
          end
       | _ -> Lwt.return_none
     in match id with
        | Some _id -> Lwt.return portal
        | None -> Lwt.fail_with "Invalid stream opening server-side."

let starttls _ = Lwt.fail_with "STARTTLS is unimplemented in WebSockets."

let _encrypted _ = true