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

type step =
  | Feature of Feature.requirement * features
  | Logged_in of features

type state =
  | Disconnected
  | Connecting of Portal.domain
  (* TCP/WebSocket connected. *)
  | Connected of Portal.socket
  (* Stream negotiation *)
  | Negotiating of Portal.t * step

(** An XMPP session. This type contains a signal representing the state of an XMPP
    connection, and its update function. *)
type t = {
    state : state signal;
    update : state -> unit;
  }

(** [create domain config] is a promise containing a session representing the portal
    connected to the XMPP server located at [domain].

    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 (config : config) : t Lwt.t =
  let features_next_state = function
    | feature :: rest -> Feature (feature, rest)
    (* No features returned by stream start. Connection is completed. *)
    | [] -> Logged_in []
  and eq s1 s2 =  (* TODO: move this closer to the state type *)
    match s1, s2 with
    | Negotiating (_, step1), Negotiating (_, step2) -> step1 = step2
    | _ -> s1 = s2
  in
  let state, update = S.create ~eq Disconnected in
  let connection_map = function
    | Connecting domain ->
       let+ socket = Portal.connect domain
       in update (Connected socket)
    | Connected socket ->
       let* portal = Portal.stream socket in
       let+ features = Stream.start portal
       in let next_state = features_next_state features
          in update (Negotiating (portal, next_state))
    | Negotiating (portal, Feature (feature, features)) ->
       let+ negotiation = negotiate feature portal config
       in let next_state = match negotiation with
            | Some socket -> Connected socket
            | None -> Negotiating (portal, features_next_state features)
          in update next_state
    | _ -> Lwt.return_unit
  in let+ () = S.map_s connection_map state >|= S.keep;
     in { state; update }

let connect (session : t) (domain : Portal.domain) : unit =
  session.update (Connecting domain)