// SPDX-License-Identifier: MIT OR Apache-2.0

//! API for managing members and their key bundles.
use std::fmt::Debug;
use std::marker::PhantomData;

use p2panda_auth::traits::Conditions;
use p2panda_core::VerifyingKey;
use p2panda_encryption::key_bundle::{KeyBundleError, Lifetime, LongTermKeyBundle};
use p2panda_encryption::key_manager::{KeyManager, KeyManagerError, KeyManagerState};
use p2panda_encryption::key_registry::{KeyRegistry, KeyRegistryError, KeyRegistryState};
use p2panda_encryption::traits::{KeyBundle, PreKeyManager};
use p2panda_encryption::{Rng, RngError};
use p2panda_store::Transaction;
use p2panda_store::key_registry::KeyRegistryStore;
use p2panda_store::key_secrets::KeySecretsStore;
use thiserror::Error;

use crate::event::Event;
use crate::forge::Forge;
use crate::manager::StoreError;
use crate::member::Member;
use crate::message::SpacesArgs;
use crate::{Config, Credentials, MemberId};

/// Manager for functionality relating to a peers identity, holds all cryptographic secrets for key
/// agreement and signatures.
///
/// Exposes an API for publishing and storing/retrieving key bundles, including rotating our own
/// when they expire, as well as methods for "forging" (constructing and signing messages) which
/// are signed with the peers private key.
///
/// **Warning:** Neither of a peers keys should be rotated individually, this would result in
/// undefined behavior. Rotating both keys is possible but will result in the loss of access to
/// existing spaces.
#[derive(Debug)]
pub struct IdentityManager<S, F, C> {
    key_store: S,
    forge: F,
    credentials: Credentials,
    config: Config,
    rng: Rng,
    _marker: PhantomData<C>,
}

impl<S, F, C> IdentityManager<S, F, C>
where
    S: KeyRegistryStore + KeySecretsStore + Transaction,
    F: Forge<C>,
    C: Conditions,
{
    #[allow(clippy::result_large_err)]
    pub fn new(
        key_store: S,
        forge: F,
        credentials: Credentials,
        config: Config,
        rng: &Rng,
    ) -> Result<Self, IdentityError<F, C>> {
        Ok(Self {
            key_store,
            forge,
            credentials,
            config,
            rng: Rng::from_rng(rng)?,
            _marker: PhantomData,
        })
    }

    /// The public key of the local actor.
    pub(crate) fn id(&self) -> MemberId {
        self.credentials.verifying_key()
    }

    /// The local actor id and their long-term key bundle.
    ///
    /// Note: Key bundle will be rotated if the latest is reaching it's configured expiry date.
    pub(crate) async fn me(&self) -> Result<Member, IdentityError<F, C>> {
        Ok(Member::new(self.id(), self.key_bundle().await?))
    }

    /// Returns "latest", publishable key bundle of us or automatically generates a new one if
    /// either nothing was generated yet, if previous bundles expired or are about to be expired
    /// (given an additional "pessimistic" rotation window).
    async fn key_bundle(&self) -> Result<LongTermKeyBundle, IdentityError<F, C>> {
        let key_manager_y = self.key_manager().await?;

        let valid_bundle = match KeyManager::prekey_bundle(&key_manager_y) {
            Ok(bundle) => bundle
                .lifetime()
                .verify_with_window(self.config.pre_key_rotate_after)
                .map_or(None, |_| Some(bundle)),
            Err(KeyManagerError::NoPreKeysAvailable) => None,
            Err(err) => return Err(err.into()),
        };

        if let Some(bundle) = valid_bundle {
            return Ok(bundle);
        }

        // Automatically rotate pre key.
        let key_manager_y_i = KeyManager::rotate_prekey(
            key_manager_y,
            Lifetime::new(self.config.pre_key_lifetime.as_secs()),
            &self.rng,
        )?;

        let key_registry_y = self.key_registry().await?;

        // Register our own key bundle.
        let key_bundle = KeyManager::prekey_bundle(&key_manager_y_i)?;
        let key_registry_y_i =
            KeyRegistry::add_longterm_bundle(key_registry_y, self.id(), key_bundle.clone())?;

        // Clean up expired key bundles ("garbage collection").
        let key_manager_y_ii = KeyManager::remove_expired(key_manager_y_i);
        let key_registry_y_ii = KeyRegistry::remove_expired(key_registry_y_i);

        // Persist new state in store.
        {
            let permit = self
                .key_store
                .begin()
                .await
                .map_err(|err| StoreError::Transaction(err.to_string()))?;

            self.key_store
                .set_prekey_secrets(key_manager_y_ii.prekey_bundles())
                .await
                .map_err(|err| StoreError::KeySecretStore(err.to_string()))?;
            self.key_store
                .set_key_registry(&key_registry_y_ii)
                .await
                .map_err(|err| StoreError::KeyRegistryStore(err.to_string()))?;

            self.key_store
                .commit(permit)
                .await
                .map_err(|err| StoreError::Transaction(err.to_string()))?;
        }

        Ok(key_bundle)
    }

    /// Returns `true` if my latest key bundle has expired or is about to expire.
    pub async fn key_bundle_expired(&self) -> Result<bool, IdentityError<F, C>> {
        let key_manager_y = self.key_manager().await?;
        match KeyManager::prekey_bundle(&key_manager_y) {
            Ok(bundle) => Ok(bundle
                .lifetime()
                .verify_with_window(self.config.pre_key_rotate_after)
                .is_err()),
            Err(KeyManagerError::NoPreKeysAvailable) => Ok(true),
            Err(err) => Err(err.into()),
        }
    }

    /// Forge a key bundle message containing my latest key bundle.
    ///
    /// Note: Key bundle will be rotated if the latest is reaching it's configured expiry date.
    pub async fn key_bundle_message(&mut self) -> Result<F::Message, IdentityError<F, C>> {
        let args = SpacesArgs::KeyBundle {
            key_bundle: self.key_bundle().await?,
        };
        let message = self.forge.forge(args).await.map_err(IdentityError::Forge)?;
        Ok(message)
    }

    /// Register a member with long-term key bundle material.
    ///
    /// Throws an error if provided key bundle has an invalid signature or expired.
    //
    // NOTE: **Security:** This method does _only_ validate if the pre-key signature maps to the
    // given identity key but **not** if the member's handle / id is authentic. Applications need to
    // provide an authentication scheme and validate `Member` before calling this method to prevent
    // impersonation attacks.
    pub async fn register_member(&mut self, member: &Member) -> Result<(), IdentityError<F, C>> {
        let pki = {
            let y = self.key_registry().await?;
            KeyRegistry::add_longterm_bundle(y, member.id(), member.key_bundle().clone())?
        };

        {
            let permit = self
                .key_store
                .begin()
                .await
                .map_err(|err| StoreError::Transaction(err.to_string()))?;

            self.key_store
                .set_key_registry(&pki)
                .await
                .map_err(|err| StoreError::KeyRegistryStore(err.to_string()))?;

            self.key_store
                .commit(permit)
                .await
                .map_err(|err| StoreError::Transaction(err.to_string()))?;
        }

        Ok(())
    }

    /// Process a key bundle received from the network.
    pub async fn process_key_bundle(
        &mut self,
        author: MemberId,
        key_bundle: &LongTermKeyBundle,
    ) -> Result<Event<C>, IdentityError<F, C>> {
        key_bundle.verify()?;
        let member = Member::new(author, key_bundle.clone());
        self.register_member(&member).await?;
        Ok(Event::KeyBundle { author })
    }

    pub async fn forge(&mut self, args: SpacesArgs<C>) -> Result<F::Message, IdentityError<F, C>> {
        self.forge.forge(args).await.map_err(IdentityError::Forge)
    }

    /// Assemble and return key manager state from persisted pre-key bundles and identity secret.
    pub async fn key_manager(&self) -> Result<KeyManagerState, IdentityError<F, C>> {
        let y = self
            .key_store
            .get_prekey_secrets()
            .await
            .map_err(|err| StoreError::KeySecretStore(err.to_string()))?
            .unwrap_or_default();

        Ok(KeyManager::init_from_prekey_bundles(
            &self.credentials.identity_secret(),
            y,
        )?)
    }

    pub async fn key_registry(&self) -> Result<KeyRegistryState<MemberId>, IdentityError<F, C>> {
        let y = match self
            .key_store
            .get_key_registry()
            .await
            .map_err(|err| StoreError::KeyRegistryStore(err.to_string()))?
        {
            Some(y) => y,
            None => KeyRegistry::init(),
        };

        Ok(y)
    }
}

#[derive(Debug, Error)]
#[allow(clippy::large_enum_variant)]
pub enum IdentityError<F, C>
where
    F: Forge<C>,
    C: Conditions,
{
    #[error("{0}")]
    Forge(F::Error),

    #[error(transparent)]
    KeyManager(#[from] KeyManagerError),

    #[error(transparent)]
    KeyRegistry(#[from] KeyRegistryError),

    #[error(transparent)]
    Rng(#[from] RngError),

    #[error(transparent)]
    KeyBundle(#[from] KeyBundleError),

    #[error("received long-term key bundle for {0} on message signed by unexpected author {1}")]
    KeyBundleAuthor(VerifyingKey, VerifyingKey),

    #[error(transparent)]
    Store(#[from] StoreError),
}

#[cfg(test)]
mod tests {
    use std::borrow::Borrow;
    use std::time::Duration;

    use p2panda_core::traits::Provenance;
    use p2panda_encryption::Rng;
    use p2panda_encryption::key_bundle::LongTermKeyBundle;
    use p2panda_encryption::key_registry::KeyRegistry;
    use p2panda_encryption::traits::{KeyBundle, PreKeyRegistry};
    use p2panda_store::SqliteStore;

    use crate::message::SpacesArgs;
    use crate::test_utils::TestForge;
    use crate::{Config, Credentials};

    use super::IdentityManager;

    #[tokio::test]
    async fn me_returns_valid_member() {
        let rng = Rng::from_seed([1; 32]);

        let credentials = Credentials::from_rng(&rng).unwrap();
        let config = Config::default();

        let store = SqliteStore::temporary().await;
        let forge = TestForge::new(store.clone(), credentials.signing_key());

        let identity_manager =
            IdentityManager::new(store, forge, credentials.clone(), config, &rng).unwrap();

        let me = identity_manager.me().await.unwrap();
        let bundle: &LongTermKeyBundle = me.key_bundle();
        let actor_id = credentials.verifying_key();

        assert_eq!(me.id(), actor_id);
        assert!(bundle.verify().is_ok());
    }

    #[tokio::test]
    async fn key_bundle_message_forged() {
        let rng = Rng::from_seed([1; 32]);

        let credentials = Credentials::from_rng(&rng).unwrap();
        let config = Config::default();

        let store = SqliteStore::temporary().await;
        let forge = TestForge::new(store.clone(), credentials.signing_key());

        let mut identity_manager =
            IdentityManager::new(store, forge, credentials.clone(), config, &rng).unwrap();

        let msg = identity_manager.key_bundle_message().await.unwrap();

        let actor_id = credentials.signing_key().verifying_key();
        assert_eq!(msg.author(), actor_id);
        match msg.borrow() {
            SpacesArgs::KeyBundle { key_bundle } => {
                assert!(key_bundle.verify().is_ok());
            }
            _ => panic!("expected key bundle message"),
        }
    }

    #[tokio::test]
    async fn process_key_bundle_registers_member() {
        let alice_rng = Rng::from_seed([1; 32]);
        let alice_credentials = Credentials::from_rng(&alice_rng).unwrap();
        let alice_config = Config::default();

        let alice_store = SqliteStore::temporary().await;
        let alice_forge = TestForge::new(alice_store.clone(), alice_credentials.signing_key());

        let mut alice_identity_manager = IdentityManager::new(
            alice_store,
            alice_forge,
            alice_credentials,
            alice_config,
            &alice_rng,
        )
        .unwrap();

        let bob_rng = Rng::from_seed([2; 32]);
        let bob_credentials = Credentials::from_rng(&bob_rng).unwrap();
        let bob_config = Config::default();

        let bob_store = SqliteStore::temporary().await;
        let bob_forge = TestForge::new(bob_store.clone(), bob_credentials.signing_key());

        let bob_identity_manager = IdentityManager::new(
            bob_store,
            bob_forge,
            bob_credentials.clone(),
            bob_config,
            &bob_rng,
        )
        .unwrap();
        let bob_id = bob_credentials.verifying_key().into();

        let bob_member = bob_identity_manager.me().await.unwrap();
        let bob_bundle = bob_member.key_bundle();
        alice_identity_manager
            .process_key_bundle(bob_id, bob_bundle)
            .await
            .unwrap();

        let key_registry_y = alice_identity_manager.key_registry().await.unwrap();
        let (_, bundle): (_, Option<LongTermKeyBundle>) =
            KeyRegistry::key_bundle(key_registry_y, &bob_id).unwrap();
        assert!(bundle.is_some());
        let bundle_identity_key = bundle.unwrap().identity_key().to_owned();
        assert_eq!(bundle_identity_key, *bob_bundle.identity_key());
        assert_eq!(
            bundle_identity_key,
            bob_credentials.identity_secret().verifying_key().unwrap()
        );
    }

    #[tokio::test]
    async fn me_rotates_key_bundle_when_expired() {
        let rng = Rng::from_seed([1; 32]);
        let credentials = Credentials::from_rng(&rng).unwrap();
        let config = Config::default();

        let store = SqliteStore::temporary().await;
        let forge = TestForge::new(store.clone(), credentials.signing_key());

        let mut alice_identity_manager =
            IdentityManager::new(store, forge, credentials, config, &rng).unwrap();

        let alice_1 = alice_identity_manager.me().await.unwrap();
        let bundle_1 = alice_1.key_bundle().clone();

        // Override max. lifetime of 90 days (default) with pre-rotation window to force rotation.
        alice_identity_manager.config.pre_key_rotate_after =
            Duration::from_secs(60 * 60 * 24 * 1024);

        // Make lifetime of next key longer to "win" over the previous one, in case it is still
        // considered valid due to a race condition (both keys can be generated "at the same time").
        alice_identity_manager.config.pre_key_lifetime = Duration::from_secs(60 * 60 * 24 * 2048);

        let alice_2 = alice_identity_manager.me().await.unwrap();
        let bundle_2 = alice_2.key_bundle().clone();

        // Key bundles are valid, we only forced the generate a new one to pessimistically already
        // distribute it, but the "old" one is still fine!
        assert!(bundle_1.verify().is_ok());
        assert!(bundle_2.verify().is_ok());

        assert_ne!(
            bundle_1.signed_prekey(),
            bundle_2.signed_prekey(),
            "rotation should produce a new pre-key"
        );
    }
}
