open Lwt.Syntax
open Lwt.Infix
open Lwt_react
open Stream

type connection_step =
  | Starting_stream
  | Negotiating_feature of Feature.requirement * features

type state =
  | Disconnected
  | Opening_portal of string
  | Connecting of Portal.t * connection_step
  | Connected of Portal.t * features

(** [create domain config] is a promise containing a signal representing the portal
    connected to the XMPP server located at [domain], and all its supported features.

    When calling the function, the portal will try to connect to the provided domain.

    Basically, it conforms to {{:
    https://datatracker.ietf.org/doc/html/rfc6120#section-4.3 }}, and gets the Portal in
    a "ready" state. *)
let create (domain : string) (config : config) : (state signal * (unit -> unit)) Lwt.t =
  let needs_restart = function
    | Feature.Mechanisms _ | STARTTLS -> true
    | _ -> false
  and features_next_state portal = function
    | feature :: rest -> Connecting (portal, Negotiating_feature (feature, rest))
    (* No features returned by stream start. Connection is completed. *)
    | [] -> Connected (portal, [])
  and eq s1 s2 =  (* TODO: move this closer to the state type *)
    match s1, s2 with
    | Connecting (_, step1), Connecting (_, step2) -> step1 = step2
    | Connected (_, f1), Connected (_, f2) -> f1 = f2
    | _ -> s1 = s2
  in
  let state, update = S.create ~eq Disconnected in
  let+ () = S.map_s
              (function
               | Opening_portal domain ->
                  let+ portal = Portal.connect domain
                  in update (Connecting (portal, Starting_stream))
               | Connecting (portal, Starting_stream) ->
                  let+ features = Stream.start domain portal
                  in features_next_state portal features |> update
               | Connecting (portal, Negotiating_feature (feature, features)) ->
                  let+ () = negotiate feature portal config
                  in let next_state = if needs_restart (Feature.unwrap feature)
                                      then Connecting (portal, Starting_stream)
                                      else features_next_state portal features
                     in update next_state
               | _ -> Lwt.return_unit)
              state >|= S.keep;
  in state, fun () -> update (Opening_portal domain)