open Lwt.Syntax

type auth_mechanism = PLAIN  [@@deriving show { with_path = false }]

type auth_config = {
  jid : string;
  password : string;
  preferred_mechanisms : auth_mechanism list;
}

type sasl_error =
  | NotAuthorized
  | MalformedRequest

let read_sasl_error = function
  | "not-authorized" -> NotAuthorized
  | "malformed-request" -> MalformedRequest
  | _ -> failwith "Unsupported SASL error returned by the server."

type sasl_auth = (string option, sasl_error * (string * string) option) result

let send_auth_stanza (stream, push) localpart pass mechanism =
  let gen_auth = function
    | PLAIN -> Base64.encode_exn ("\x00" ^ localpart ^ "\x00" ^ pass)
  and parse_sasl_response stanza =
    let open Markup in
    let parse_additional_info = function
      | `Text t :: _ -> Some (String.concat "" t)
      | _ -> None
    and parse_descriptive_text = function
      | `Start_element ((_, "text"), [((_, "lang"), lang)]) :: `Text desc :: _ -> Some (lang, String.concat "" desc)
      | `Start_element ((_, "text"), []) :: `Text desc :: _ -> Some ("en", String.concat "" desc)
      | _ -> None
    in
    let parse_sasl_error = function
      | `Start_element ((_, error), _) :: `End_element :: rest -> (read_sasl_error error, parse_descriptive_text rest)
      | _ -> raise (Stream.InvalidStanza stanza)
    in match (string stanza |> parse_xml |> signals |> to_list) with
       | `Start_element ((_, "success"), _) :: rest -> Ok (parse_additional_info rest)
       | `Start_element ((_, "failure"), _) :: rest -> Error (parse_sasl_error rest)
       | _ -> raise (Stream.InvalidStanza stanza)
  in let xmlns = "urn:ietf:params:xml:ns:xmpp-sasl" in
     let stanza_list = [`Start_element
                          ((xmlns, "auth"),
                           [(("", "xmlns"), xmlns);
                            (("", "mechanism"), show_auth_mechanism mechanism)]);
                        `Text [gen_auth mechanism];
                        `End_element]
     in Markup.(stanza_list |> of_list |> write_xml |> to_string) |> Option.some |> push;
        let* response = Stream.get stream
        in try parse_sasl_response response |> Lwt.return
           with exn -> Lwt.fail exn

let authenticate (portal : Portal.t) (config : auth_config) =
  let {jid; password; _} = config
  (* Probably not exactly compliant with https://xmpp.org/extensions/xep-0029.html,
     but it's just for simplicity's sake in alpha. *)
  in let localpart = match String.split_on_char '@' jid with
       | [localpart; _domain] -> localpart
       | _ -> failwith "Invalid JID"
     in send_auth_stanza portal localpart password PLAIN