use aucpace::client::{AuCPaceClientPreAug, AuCPaceClientRecvServerKey};
use aucpace::server::{AuCPaceServerAugLayer, AuCPaceServerRecvClientKey};
use aucpace::{
    Client, ClientMessage, Error, PartialAugDatabase, Result, Server, ServerMessage, StrongDatabase,
};
use curve25519_dalek::{RistrettoPoint, Scalar};
use password_hash::ParamsString;
use rand_core::OsRng;
use scrypt::{Params, Scrypt};
use sha2::Sha512;

const USERNAME: &'static [u8] = b"jlpicard_1701";
const PASSWORD: &'static [u8] = b"g04tEd_c4pT41N";
const CI: &[u8] = b"test_channel_identifier";
const PRE_SSID: &[u8] = b"bestest_ssid_ever_i_promise";
const K1: usize = 16;

/// Password Verifier database which can store the info for one user
#[derive(Debug, Default)]
struct SingleUserDatabase {
    user: Option<Vec<u8>>,
    data: Option<(RistrettoPoint, Scalar, ParamsString)>,
    long_term_keypair: Option<(Scalar, RistrettoPoint)>,
}

impl StrongDatabase for SingleUserDatabase {
    type PasswordVerifier = RistrettoPoint;
    type Exponent = Scalar;

    fn lookup_verifier_strong(
        &self,
        username: &[u8],
    ) -> Option<(Self::PasswordVerifier, Self::Exponent, ParamsString)> {
        match &self.user {
            Some(stored_username) if stored_username == username => self.data.clone(),
            _ => None,
        }
    }

    fn store_verifier_strong(
        &mut self,
        username: &[u8],
        _uad: Option<&[u8]>,
        verifier: Self::PasswordVerifier,
        secret_exponent: Self::Exponent,
        params: ParamsString,
    ) {
        self.user = Some(username.to_vec());
        self.data = Some((verifier, secret_exponent, params));
    }
}

impl PartialAugDatabase for SingleUserDatabase {
    type PrivateKey = Scalar;
    type PublicKey = RistrettoPoint;

    fn lookup_long_term_keypair(
        &self,
        username: &[u8],
    ) -> Option<(Self::PrivateKey, Self::PublicKey)> {
        match &self.user {
            Some(stored_user) if stored_user == username => self.long_term_keypair.clone(),
            _ => None,
        }
    }

    fn store_long_term_keypair(
        &mut self,
        username: &[u8],
        priv_key: Self::PrivateKey,
        pub_key: Self::PublicKey,
    ) -> Result<()> {
        match &self.user {
            Some(stored_user) if stored_user == username => {
                self.long_term_keypair = Some((priv_key, pub_key));
                Ok(())
            }
            _ => Err(Error::UserNotRegistered),
        }
    }
}

/// test the default key agreement - establish SSID, normal explicit mutual auth, non-strong version
#[test]
fn test_key_agreement() -> Result<()> {
    let (mut base_client, mut base_server, database) = init()?;

    // ===== SSID Establishment =====
    let (server, server_message) = base_server.begin();
    let (client, client_message) = base_client.begin();

    // server receives client nonce
    let server = if let ClientMessage::Nonce(client_nonce) = client_message {
        server.agree_ssid(client_nonce)
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // client receives server nonce
    let client = if let ServerMessage::Nonce(server_nonce) = server_message {
        client.agree_ssid(server_nonce)
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // do the middle bits
    let (client, server, client_message, server_message) = test_core(client, server, &database)?;

    // Server receives client pub key
    let server = if let ClientMessage::PublicKey(client_pubkey) = client_message {
        server.receive_client_pubkey(client_pubkey)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // Client receives server pub key and generate the client's authenticators
    let (client, client_message) = if let ServerMessage::PublicKey(server_pubkey) = server_message {
        client.receive_server_pubkey(server_pubkey)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // ===== Explicit Mutual Authentication =====
    let (server_key, server_message) = if let ClientMessage::Authenticator(ca) = client_message {
        server.receive_client_authenticator(ca)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    let client_key = if let ServerMessage::Authenticator(sa) = server_message {
        client.receive_server_authenticator(sa)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // assert that both threads arrived at the same key
    assert_eq!(client_key, server_key);

    Ok(())
}

/// test the default key agreement - establish SSID, normal explicit mutual auth, non-strong version
#[test]
fn test_key_agreement_implicit_auth() -> Result<()> {
    let (mut base_client, mut base_server, database) = init()?;

    // ===== SSID Establishment =====
    let (server, server_message) = base_server.begin();
    let (client, client_message) = base_client.begin();

    // server receives client nonce
    let server = if let ClientMessage::Nonce(client_nonce) = client_message {
        server.agree_ssid(client_nonce)
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // client receives server nonce
    let client = if let ServerMessage::Nonce(server_nonce) = server_message {
        client.agree_ssid(server_nonce)
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // do the middle bits
    let (client, server, client_message, server_message) = test_core(client, server, &database)?;

    // Server receives client pub key
    let server_key = if let ClientMessage::PublicKey(client_pubkey) = client_message {
        server.implicit_auth(client_pubkey)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // Client receives server pub key and generate the client's authenticators
    let client_key = if let ServerMessage::PublicKey(server_pubkey) = server_message {
        client.implicit_auth(server_pubkey)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // assert that both threads arrived at the same key
    assert_eq!(client_key, server_key);

    Ok(())
}

/// test the default key agreement - establish SSID, normal explicit mutual auth, non-strong version
#[test]
fn test_key_agreement_prestablished_ssid() -> Result<()> {
    let (mut base_client, mut base_server, database) = init()?;

    // ===== SSID Establishment =====
    let server = base_server.begin_prestablished_ssid(PRE_SSID)?;
    let client = base_client.begin_prestablished_ssid(PRE_SSID)?;

    // do the middle bits
    let (client, server, client_message, server_message) = test_core(client, server, &database)?;

    // Server receives client pub key
    let server = if let ClientMessage::PublicKey(client_pubkey) = client_message {
        server.receive_client_pubkey(client_pubkey)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // Client receives server pub key and generate the client's authenticators
    let (client, client_message) = if let ServerMessage::PublicKey(server_pubkey) = server_message {
        client.receive_server_pubkey(server_pubkey)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // ===== Explicit Mutual Authentication =====
    let (server_key, server_message) = if let ClientMessage::Authenticator(ca) = client_message {
        server.receive_client_authenticator(ca)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    let client_key = if let ServerMessage::Authenticator(sa) = server_message {
        client.receive_server_authenticator(sa)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // assert that both threads arrived at the same key
    assert_eq!(client_key, server_key);

    Ok(())
}

/// test the default key agreement - establish SSID, normal explicit mutual auth, non-strong version
#[test]
fn test_key_agreement_prestablished_ssid_implicit_auth() -> Result<()> {
    let (mut base_client, mut base_server, database) = init()?;

    // ===== SSID Establishment =====
    let server = base_server.begin_prestablished_ssid(PRE_SSID)?;
    let client = base_client.begin_prestablished_ssid(PRE_SSID)?;

    // do the middle bits
    let (client, server, client_message, server_message) = test_core(client, server, &database)?;

    // Server receives client pub key
    let server_key = if let ClientMessage::PublicKey(client_pubkey) = client_message {
        server.implicit_auth(client_pubkey)?
    } else {
        panic!("Received invalid client message {:?}", client_message);
    };

    // Client receives server pub key and generate the client's authenticators
    let client_key = if let ServerMessage::PublicKey(server_pubkey) = server_message {
        client.implicit_auth(server_pubkey)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // assert that both threads arrived at the same key
    assert_eq!(client_key, server_key);

    Ok(())
}

/// Perform the initialisation step for all tests
fn init() -> Result<(Client, Server, SingleUserDatabase)> {
    // Create the client, server and database
    let mut base_server = Server::new(OsRng);
    let mut base_client = Client::new(OsRng);
    let mut database: SingleUserDatabase = Default::default();

    // register a user in the database
    let params = Params::recommended();
    let registration = base_client.register_alloc_strong(USERNAME, PASSWORD, params, Scrypt)?;
    if let ClientMessage::StrongRegistration {
        username,
        secret_exponent,
        params,
        verifier,
    } = registration
    {
        database.store_verifier_strong(username, None, verifier, secret_exponent, params);
        let (private, public) = base_server.generate_long_term_keypair();
        database.store_long_term_keypair(username, private, public)?;
    }

    Ok((base_client, base_server, database))
}

/// perform all the middle steps common to all tests
fn test_core(
    client: AuCPaceClientPreAug<Sha512, Scrypt, K1>,
    server: AuCPaceServerAugLayer<Sha512, K1>,
    database: &SingleUserDatabase,
) -> Result<(
    AuCPaceClientRecvServerKey<Sha512, K1>,
    AuCPaceServerRecvClientKey<Sha512, K1>,
    ClientMessage<K1>,
    ServerMessage<K1>,
)> {
    // ===== Augmentation Layer =====
    // client initiates the augmentation phase
    let (client, client_message) = client.start_augmentation_strong(USERNAME, PASSWORD, &mut OsRng);

    // server generates augmentation info from client's username
    let (server, server_message) =
        if let ClientMessage::StrongUsername { username, blinded } = client_message {
            server.generate_client_info_partial_strong(username, blinded, database, OsRng)?
        } else {
            panic!("Received invalid client message {:?}", client_message);
        };

    // client receives the info and moves into the CPace step
    let client = if let ServerMessage::StrongAugmentationInfo {
        x_pub,
        blinded_salt,
        pbkdf_params,
        ..
    } = server_message
    {
        let params = {
            // its christmas time!
            let log_n = pbkdf_params.get_str("ln").unwrap().parse().unwrap();
            let r = pbkdf_params.get_str("r").unwrap().parse().unwrap();
            let p = pbkdf_params.get_str("p").unwrap().parse().unwrap();

            Params::new(log_n, r, p, Params::RECOMMENDED_LEN).unwrap()
        };
        client.generate_cpace_alloc(x_pub, blinded_salt, params, Scrypt)?
    } else {
        panic!("Received invalid server message {:?}", server_message);
    };

    // ===== CPace substep =====
    let (server, server_message) = server.generate_public_key(CI);
    let (client, client_message) = client.generate_public_key(CI, &mut OsRng);

    Ok((client, server, client_message, server_message))
}