(** Authentication operations. *)

let src = Logs.Src.create "matrix.auth" ~doc:"Matrix authentication"
module Log = (val Logs.src_log src : Logs.LOG)

(* Login flow types *)
type login_flow =
  | Password
  | Token
  | Sso
  | Unknown of string

let login_flow_of_string = function
  | "m.login.password" -> Password
  | "m.login.token" -> Token
  | "m.login.sso" -> Sso
  | s -> Unknown s

let login_flow_to_string = function
  | Password -> "m.login.password"
  | Token -> "m.login.token"
  | Sso -> "m.login.sso"
  | Unknown s -> s

(* JSON codecs for login flows response *)
let login_flow_jsont =
  Jsont.of_of_string ~kind:"login_flow"
    ~enc:login_flow_to_string
    (fun s -> Ok (login_flow_of_string s))

let login_flow_obj_jsont =
  Jsont.Object.map
    ~kind:"login_flow_object"
    (fun flow_type -> flow_type)
  |> Jsont.Object.mem "type" login_flow_jsont
  |> Jsont.Object.finish

let login_flows_response_jsont =
  Jsont.Object.map
    ~kind:"login_flows_response"
    (fun flows -> flows)
  |> Jsont.Object.mem "flows" (Jsont.list login_flow_obj_jsont)
  |> Jsont.Object.finish

let get_login_flows client =
  match Client.get client ~path:"/login" () with
  | Error e -> Error e
  | Ok body -> Client.decode_response login_flows_response_jsont body

(* Login parameters *)
type login_params = {
  device_id : string option;
  initial_device_display_name : string option;
}

let default_login_params = {
  device_id = None;
  initial_device_display_name = None;
}

(* Login request codec - write-only types for JSON encoding *)
type login_request = {
  req_type : string;
  identifier : login_identifier;
  password : string option;
  token : string option;
  device_id : string option;
  initial_device_display_name : string option;
} [@@warning "-69"]

and login_identifier = {
  id_type : string;
  user : string option;
} [@@warning "-69"]

let login_identifier_jsont =
  Jsont.Object.(
    map ~kind:"login_identifier"
      (fun id_type user -> { id_type; user })
    |> mem "type" Jsont.string ~enc:(fun t -> t.id_type)
    |> opt_mem "user" Jsont.string ~enc:(fun t -> t.user)
    |> finish)

let login_request_jsont =
  Jsont.Object.(
    map ~kind:"login_request"
      (fun req_type identifier password token device_id initial_device_display_name ->
        { req_type; identifier; password; token; device_id; initial_device_display_name })
    |> mem "type" Jsont.string ~enc:(fun t -> t.req_type)
    |> mem "identifier" login_identifier_jsont ~enc:(fun t -> t.identifier)
    |> opt_mem "password" Jsont.string ~enc:(fun t -> t.password)
    |> opt_mem "token" Jsont.string ~enc:(fun t -> t.token)
    |> opt_mem "device_id" Jsont.string ~enc:(fun t -> t.device_id)
    |> opt_mem "initial_device_display_name" Jsont.string
         ~enc:(fun t -> t.initial_device_display_name)
    |> finish)

(* Login response codec *)
type login_response = {
  user_id : Matrix_proto.Id.User_id.t;
  access_token : string;
  device_id : Matrix_proto.Id.Device_id.t;
  refresh_token : string option;
}

let login_response_jsont =
  Jsont.Object.map
    ~kind:"login_response"
    (fun user_id access_token device_id refresh_token ->
      { user_id; access_token; device_id; refresh_token })
  |> Jsont.Object.mem "user_id" Matrix_proto.Id.User_id.jsont
  |> Jsont.Object.mem "access_token" Jsont.string
  |> Jsont.Object.mem "device_id" Matrix_proto.Id.Device_id.jsont
  |> Jsont.Object.opt_mem "refresh_token" Jsont.string
  |> Jsont.Object.finish

let response_to_session resp : Client.session =
  { user_id = resp.user_id;
    access_token = resp.access_token;
    device_id = resp.device_id;
    refresh_token = resp.refresh_token;
  }

let login_password client ~user ~password ?(params = default_login_params) () =
  Log.info (fun m -> m "Logging in as %s" user);
  let request = {
    req_type = "m.login.password";
    identifier = { id_type = "m.id.user"; user = Some user };
    password = Some password;
    token = None;
    device_id = params.device_id;
    initial_device_display_name = params.initial_device_display_name;
  } in
  match Client.encode_body login_request_jsont request with
  | Error e -> Error e
  | Ok body ->
      match Client.post_unauthenticated client ~path:"/login" ~body () with
      | Error e ->
          Log.err (fun m -> m "Login failed for user %s" user);
          Error e
      | Ok body ->
          match Client.decode_response login_response_jsont body with
          | Error e -> Error e
          | Ok resp ->
              Log.info (fun m -> m "Login successful: user_id=%s device_id=%s"
                (Matrix_proto.Id.User_id.to_string resp.user_id)
                (Matrix_proto.Id.Device_id.to_string resp.device_id));
              Ok (response_to_session resp)

let login_token client ~token ?(params = default_login_params) () =
  let request = {
    req_type = "m.login.token";
    identifier = { id_type = "m.id.user"; user = None };
    password = None;
    token = Some token;
    device_id = params.device_id;
    initial_device_display_name = params.initial_device_display_name;
  } in
  match Client.encode_body login_request_jsont request with
  | Error e -> Error e
  | Ok body ->
      match Client.post_unauthenticated client ~path:"/login" ~body () with
      | Error e -> Error e
      | Ok body ->
          match Client.decode_response login_response_jsont body with
          | Error e -> Error e
          | Ok resp -> Ok (response_to_session resp)

(* Token refresh *)
type refresh_request = {
  refresh_token : string;
} [@@warning "-69"]

let refresh_request_jsont =
  Jsont.Object.(
    map ~kind:"refresh_request"
      (fun refresh_token -> { refresh_token })
    |> mem "refresh_token" Jsont.string ~enc:(fun t -> t.refresh_token)
    |> finish)

type refresh_response = {
  access_token : string;
  refresh_token : string option;
}

let refresh_response_jsont =
  Jsont.Object.map
    ~kind:"refresh_response"
    (fun access_token refresh_token -> { access_token; refresh_token })
  |> Jsont.Object.mem "access_token" Jsont.string
  |> Jsont.Object.opt_mem "refresh_token" Jsont.string
  |> Jsont.Object.finish

let refresh_token client ~refresh_token =
  let request = { refresh_token } in
  match Client.encode_body refresh_request_jsont request with
  | Error e -> Error e
  | Ok body ->
      match Client.post_unauthenticated client ~path:"/refresh" ~body () with
      | Error e -> Error e
      | Ok body ->
          match Client.decode_response refresh_response_jsont body with
          | Error e -> Error e
          | Ok resp -> Ok (resp.access_token, resp.refresh_token)

(* Logout *)
let logout client =
  match Client.post client ~path:"/logout" ~body:"{}" () with
  | Error e -> Error e
  | Ok _ -> Ok ()

let logout_all client =
  match Client.post client ~path:"/logout/all" ~body:"{}" () with
  | Error e -> Error e
  | Ok _ -> Ok ()

(* Registration *)
type registration_kind =
  | User
  | Guest

type register_request = {
  kind : string option;
  username : string option;
  password : string option;
  device_id : string option;
  initial_device_display_name : string option;
  inhibit_login : bool option;
} [@@warning "-69"]

let register_request_jsont =
  Jsont.Object.(
    map ~kind:"register_request"
      (fun username password device_id initial_device_display_name inhibit_login ->
        { kind = None; username; password; device_id; initial_device_display_name; inhibit_login })
    |> opt_mem "username" Jsont.string ~enc:(fun t -> t.username)
    |> opt_mem "password" Jsont.string ~enc:(fun t -> t.password)
    |> opt_mem "device_id" Jsont.string ~enc:(fun t -> t.device_id)
    |> opt_mem "initial_device_display_name" Jsont.string
         ~enc:(fun t -> t.initial_device_display_name)
    |> opt_mem "inhibit_login" Jsont.bool ~enc:(fun t -> t.inhibit_login)
    |> finish)

type register_response = {
  user_id : Matrix_proto.Id.User_id.t;
  access_token : string option;
  device_id : string option;
  refresh_token : string option;
}

let register_response_jsont =
  Jsont.Object.map
    ~kind:"register_response"
    (fun user_id access_token device_id refresh_token ->
      { user_id; access_token; device_id; refresh_token })
  |> Jsont.Object.mem "user_id" Matrix_proto.Id.User_id.jsont
  |> Jsont.Object.opt_mem "access_token" Jsont.string
  |> Jsont.Object.opt_mem "device_id" Jsont.string
  |> Jsont.Object.opt_mem "refresh_token" Jsont.string
  |> Jsont.Object.finish

let register client ?kind ?username ?password ?device_id ?initial_device_display_name ?inhibit_login () =
  let kind_str = match kind with
    | Some Guest -> Some "guest"
    | Some User | None -> None
  in
  let query = match kind_str with
    | Some k -> Some [("kind", k)]
    | None -> None
  in
  let request = {
    kind = None;
    username;
    password;
    device_id;
    initial_device_display_name;
    inhibit_login;
  } in
  match Client.encode_body register_request_jsont request with
  | Error e -> Error e
  | Ok body ->
      match Client.post_unauthenticated client ~path:"/register" ?query ~body () with
      | Error e -> Error e
      | Ok body ->
          match Client.decode_response register_response_jsont body with
          | Error e -> Error e
          | Ok resp ->
              match resp.access_token, resp.device_id with
              | Some access_token, Some device_id ->
                  let device_id = Matrix_proto.Id.Device_id.of_string_exn device_id in
                  Ok { Client.user_id = resp.user_id;
                       access_token;
                       device_id;
                       refresh_token = resp.refresh_token }
              | _ ->
                  Error (Error.Json_error "Registration succeeded but no session returned (inhibit_login may be true)")

(* Whoami *)
type whoami_response = {
  user_id : Matrix_proto.Id.User_id.t;
}

let whoami_response_jsont =
  Jsont.Object.map
    ~kind:"whoami_response"
    (fun user_id -> { user_id })
  |> Jsont.Object.mem "user_id" Matrix_proto.Id.User_id.jsont
  |> Jsont.Object.finish

let whoami client =
  match Client.get client ~path:"/account/whoami" () with
  | Error e -> Error e
  | Ok body ->
      match Client.decode_response whoami_response_jsont body with
      | Error e -> Error e
      | Ok resp -> Ok resp.user_id