macltok

use std::{

    collections::HashMap,

    path::{Path, PathBuf},

};

use matrix_sdk::{

    Client, ClientBuilder, ThreadingSupport,

    authentication::{

        AuthSession,

        matrix::MatrixSession,

        oauth::{

            ClientId, ClientRegistrationData, CsrfToken, OAuthSession, UserSession,

            registration::{ApplicationType, ClientMetadata, Localized, OAuthGrantType},

        },

    },

    encryption::{BackupDownloadStrategy, EncryptionSettings},

    ruma::{OwnedDeviceId, serde::Raw},

};

use serde::{Deserialize, Serialize};

use thiserror::Error;

use url::Url;

use crate::{ClientHandle, Engine, EngineConfig, EngineError, EventStream};

pub struct Session {

    client: Client,

    record: SessionRecord,

}

pub struct PasswordLogin {

    pub homeserver: String,

    pub username: String,

    pub password: String,

    pub device_display_name: Option<String>,

    /// set after a soft logout so the store's keys still belong to us

    pub device_id: Option<OwnedDeviceId>,

}

pub struct OAuthLogin {

    pub homeserver: String,

    pub redirect_uri: String,

    pub client_uri: String,

    pub client_name: Option<String>,

    pub client_id: Option<String>,

    pub device_id: Option<OwnedDeviceId>,

    pub application_type: OAuthApplicationType,

}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]

pub enum OAuthApplicationType {

    Native,

    Web,

}

pub struct PendingOAuthLogin {

    client: Client,

    authorization_url: String,

    state: Option<CsrfToken>,

}

impl PendingOAuthLogin {

    #[must_use]

    pub fn authorization_url(&self) -> &str {

        &self.authorization_url

    }

    #[must_use]

    pub fn registered_client_id(&self) -> String {

        self.client

            .oauth()

            .client_id()

            .expect("OAuth client was registered before authorization")

            .as_str()

            .to_owned()

    }

    #[must_use]

    pub const fn is_pending(&self) -> bool {

        self.state.is_some()

    }

    pub async fn finish(&mut self, callback_url: &str) -> Result<Session, AccountError> {

        let callback_url = Url::parse(callback_url)

            .map_err(|error| account_error("parsing the OAuth callback URL", error))?;

        let state = self

            .state

            .take()

            .ok_or(AccountError::OAuthCallbackAlreadySubmitted)?;

        let oauth = self.client.oauth();

        if let Err(error) = oauth.finish_login(callback_url.into()).await {

            oauth.abort_login(&state).await;

            return Err(account_error("OAuth login", error));

        }

        Session::from_client(self.client.clone())

    }

    pub async fn abort(mut self) {

        if let Some(state) = self.state.take() {

            self.client.oauth().abort_login(&state).await;

        }

    }

}

pub struct StoreConfig {

    path: PathBuf,

    passphrase: String,

}

impl StoreConfig {

    #[must_use]

    pub fn new(path: impl Into<PathBuf>, passphrase: impl Into<String>) -> Self {

        Self {

            path: path.into(),

            passphrase: passphrase.into(),

        }

    }

    #[must_use]

    pub fn path(&self) -> &Path {

        &self.path

    }

}

/// Serializable authentication secrets. Store this separately from the

/// encrypted SQLite store in protected, preferably encrypted storage.

/// A saved record must be replaced when `Event::SessionTokensChanged` arrives.

#[derive(Clone, Deserialize, Serialize)]

pub struct SessionRecord {

    homeserver: String,

    authentication: AuthenticationRecord,

}

#[derive(Clone, Deserialize, Serialize)]

#[serde(tag = "type", content = "session", rename_all = "snake_case")]

enum AuthenticationRecord {

    Matrix(MatrixSession),

    OAuth {

        client_id: ClientId,

        user: UserSession,

    },

}

impl SessionRecord {

    #[must_use]

    pub fn homeserver(&self) -> &str {

        &self.homeserver

    }

}

#[derive(Debug, Error)]

pub enum AccountError {

    #[error("{context} failed")]

    Operation {

        context: &'static str,

        #[source]

        source: Box<dyn std::error::Error + Send + Sync>,

    },

    #[error("the homeserver returned no authenticated session")]

    MissingSession,

    #[error("the SDK returned an unsupported authentication session")]

    UnsupportedAuthentication,

    #[error("the OAuth callback was already submitted")]

    OAuthCallbackAlreadySubmitted,

    #[error(transparent)]

    Engine(#[from] EngineError),

}

impl Session {

    pub async fn login_password(

        login: PasswordLogin,

        store: StoreConfig,

    ) -> Result<Self, AccountError> {

        let client = matrix_client_builder(&store)

            .server_name_or_homeserver_url(login.homeserver.trim())

            .build()

            .await

            .map_err(|error| account_error("building the Matrix client", error))?;

        let mut request = client

            .matrix_auth()

            .login_username(login.username.trim(), &login.password)

            .request_refresh_token();

        if let Some(name) = login.device_display_name.as_deref() {

            request = request.initial_device_display_name(name);

        }

        if let Some(device_id) = login.device_id.as_deref() {

            request = request.device_id(device_id.as_str());

        }

        request

            .send()

            .await

            .map_err(|error| account_error("password login", error))?;

        Self::from_client(client)

    }

    pub async fn begin_oauth(

        login: OAuthLogin,

        store: StoreConfig,

    ) -> Result<PendingOAuthLogin, AccountError> {

        let redirect_uri = Url::parse(login.redirect_uri.trim())

            .map_err(|error| account_error("parsing the OAuth redirect URI", error))?;

        let client_uri = Url::parse(login.client_uri.trim())

            .map_err(|error| account_error("parsing the OAuth client URI", error))?;

        let client = matrix_client_builder(&store)

            .server_name_or_homeserver_url(login.homeserver.trim())

            .build()

            .await

            .map_err(|error| account_error("building the Matrix client", error))?;

        let application_type = match login.application_type {

            OAuthApplicationType::Native => ApplicationType::Native,

            OAuthApplicationType::Web => ApplicationType::Web,

        };

        let mut metadata = ClientMetadata::new(

            application_type,

            vec![OAuthGrantType::AuthorizationCode {

                redirect_uris: vec![redirect_uri.clone()],

            }],

            Localized::new(client_uri, []),

        );

        metadata.client_name = login.client_name.map(|name| Localized::new(name, []));

        let metadata = Raw::new(&metadata)

            .map_err(|error| account_error("serializing OAuth client metadata", error))?;

        let mut registration = ClientRegistrationData::new(metadata);

        if let Some(client_id) = login.client_id {

            registration.static_registrations = Some(HashMap::from([(

                client.homeserver(),

                ClientId::new(client_id),

            )]));

        }

        let authorization = client

            .oauth()

            .login(redirect_uri, login.device_id, Some(registration), None)

            .build()

            .await

            .map_err(|error| account_error("starting OAuth login", error))?;

        Ok(PendingOAuthLogin {

            client,

            authorization_url: authorization.url.to_string(),

            state: Some(authorization.state),

        })

    }

    /// Restore a saved record with the same store location and passphrase.

    /// The caller must also handle refresh, rejection, and explicit logout.

    pub async fn restore(record: SessionRecord, store: StoreConfig) -> Result<Self, AccountError> {

        let client = matrix_client_builder(&store)

            .homeserver_url(&record.homeserver)

            .build()

            .await

            .map_err(|error| account_error("reopening the Matrix store", error))?;

        let authentication = match record.authentication.clone() {

            AuthenticationRecord::Matrix(session) => AuthSession::Matrix(session),

            AuthenticationRecord::OAuth { client_id, user } => {

                OAuthSession { client_id, user }.into()

            }

        };

        client

            .restore_session(authentication)

            .await

            .map_err(|error| account_error("restoring the authenticated session", error))?;

        Ok(Self { client, record })

    }

    /// Initial login credentials only. This snapshot does not track later

    /// token rotations; use `ClientHandle::session_record` after building.

    #[must_use]

    pub fn record(&self) -> SessionRecord {

        self.record.clone()

    }

    pub async fn build(

        self,

        config: EngineConfig,

    ) -> Result<(Engine, ClientHandle, EventStream), AccountError> {

        Engine::from_client(self.client, config)

            .await

            .map_err(AccountError::from)

    }

    fn from_client(client: Client) -> Result<Self, AccountError> {

        let record = session_record(&client)?;

        Ok(Self { client, record })

    }

}

pub(crate) fn session_record(client: &Client) -> Result<SessionRecord, AccountError> {

    let authentication = match client.session().ok_or(AccountError::MissingSession)? {

        AuthSession::Matrix(session) => AuthenticationRecord::Matrix(session),

        AuthSession::OAuth(session) => {

            let OAuthSession { client_id, user } = *session;

            AuthenticationRecord::OAuth { client_id, user }

        }

        _ => return Err(AccountError::UnsupportedAuthentication),

    };

    Ok(SessionRecord {

        homeserver: client.homeserver().to_string(),

        authentication,

    })

}

fn encryption_settings() -> EncryptionSettings {

    EncryptionSettings {

        auto_enable_cross_signing: true,

        backup_download_strategy: BackupDownloadStrategy::OneShot,

        auto_enable_backups: true,

    }

}

fn matrix_client_builder(store: &StoreConfig) -> ClientBuilder {

    Client::builder()

        .with_encryption_settings(encryption_settings())

        .with_threading_support(ThreadingSupport::Enabled {

            with_subscriptions: false,

        })

        .handle_refresh_tokens()

        .sqlite_store(&store.path, Some(&store.passphrase))

}

fn account_error(

    context: &'static str,

    source: impl std::error::Error + Send + Sync + 'static,

) -> AccountError {

    AccountError::Operation {

        context,

        source: Box::new(source),

    }

}