open Lwt.Syntax

open Js_of_ocaml
let jss = Js.string
let sjs = Js.to_string

let lwt_ws (url : string) =
  (** [lwt_ws url] returns a stream and a push function for the websocket at [url].
      Pushing [None] closes the websocket.
      If the websocket is closed server-side, it's up to the caller's to close the push stream. *)
  let open Lwt_stream in
  let handle ws incoming () =
    let+ _ = iter (fun msg -> ws##send (jss msg)) incoming
    in (ws##close)
  in let stream, message = create () (* websocket -> user *)
     and incoming, push = create () (* user -> websocket *)
     in let (ws : WebSockets.webSocket Js.t) = new%js WebSockets.webSocket (jss url)
        in ws##.onmessage :=
             Dom.handler (fun x -> Some (sjs x##.data) |> message; Js._false);
           ws##.onopen :=
             Dom.handler (fun _ -> Lwt.async @@ handle ws incoming; Js._false);
           ws##.onclose :=
             Dom.handler (fun _ -> message None; Js._true);
           stream, push