use rand::RngCore;

use crate::ake::generate_kem_keypair;
use crate::envelope;
use crate::error::Result;
use crate::oprf::{BlindedElement, EvaluatedElement, OprfClient, OprfServer};
use crate::types::{
    ClientPrivateKey, ClientPublicKey, OPRF_SEED_LEN, OprfSeed, RegistrationRecord,
    RegistrationRequest, RegistrationResponse, ServerPublicKey,
};

pub struct ClientRegistrationState {
    oprf_client: OprfClient,
}

pub fn client_registration_start(
    password: &[u8],
) -> (ClientRegistrationState, RegistrationRequest) {
    let (oprf_client, blinded) = OprfClient::blind(password);

    let request = RegistrationRequest {
        blinded_element: blinded.to_bytes(),
    };

    let state = ClientRegistrationState { oprf_client };

    (state, request)
}

pub fn server_registration_respond(
    oprf_seed: &OprfSeed,
    request: &RegistrationRequest,
    server_public_key: &ServerPublicKey,
    credential_id: &[u8],
) -> Result<RegistrationResponse> {
    let blinded = BlindedElement::from_bytes(&request.blinded_element)?;

    let oprf_server = OprfServer::new(oprf_seed.clone());
    let evaluated = oprf_server.evaluate_with_credential_id(&blinded, credential_id)?;

    Ok(RegistrationResponse {
        evaluated_element: evaluated.to_bytes(),
        server_public_key: server_public_key.clone(),
    })
}

pub fn client_registration_finish(
    state: ClientRegistrationState,
    response: &RegistrationResponse,
    server_identity: Option<&[u8]>,
    client_identity: Option<&[u8]>,
) -> Result<RegistrationRecord> {
    let evaluated = EvaluatedElement::from_bytes(&response.evaluated_element)?;

    let randomized_pwd = state.oprf_client.finalize(&evaluated)?;

    let (client_kem_pk, client_kem_sk) = generate_kem_keypair();
    let client_private_key = ClientPrivateKey::new(client_kem_sk.as_bytes());
    let client_public_key = ClientPublicKey::new(client_kem_pk.as_bytes());

    let (envelope, _, _, masking_key) = envelope::store(
        &randomized_pwd,
        &response.server_public_key,
        &client_private_key,
        server_identity,
        client_identity,
    )?;

    Ok(RegistrationRecord {
        client_public_key,
        masking_key: masking_key.to_vec(),
        envelope,
    })
}

pub fn generate_oprf_seed() -> OprfSeed {
    let mut bytes = [0u8; OPRF_SEED_LEN];
    rand::thread_rng().fill_bytes(&mut bytes);
    OprfSeed::new(bytes)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ake::{generate_kem_keypair, generate_sig_keypair};

    fn create_server_keys() -> (ServerPublicKey, Vec<u8>, Vec<u8>) {
        let (kem_pk, kem_sk) = generate_kem_keypair();
        let (sig_pk, sig_sk) = generate_sig_keypair();
        let server_pk = ServerPublicKey::new(kem_pk.as_bytes(), sig_pk.as_bytes());
        (server_pk, kem_sk.as_bytes(), sig_sk.as_bytes())
    }

    #[test]
    fn test_full_registration_flow() {
        let oprf_seed = generate_oprf_seed();
        let (server_pk, _, _) = create_server_keys();
        let credential_id = b"user@example.com";
        let password = b"correct horse battery staple";

        let (client_state, request) = client_registration_start(password);

        let response =
            server_registration_respond(&oprf_seed, &request, &server_pk, credential_id).unwrap();

        let record = client_registration_finish(client_state, &response, None, None).unwrap();

        assert!(!record.client_public_key.kem_pk.is_empty());
        assert!(!record.masking_key.is_empty());
        assert!(!record.envelope.auth_tag.is_empty());
    }

    #[test]
    fn test_registration_with_identities() {
        let oprf_seed = generate_oprf_seed();
        let (server_pk, _, _) = create_server_keys();
        let credential_id = b"user@example.com";
        let password = b"password123";

        let (client_state, request) = client_registration_start(password);

        let response =
            server_registration_respond(&oprf_seed, &request, &server_pk, credential_id).unwrap();

        let record = client_registration_finish(
            client_state,
            &response,
            Some(b"server.example.com"),
            Some(b"user@example.com"),
        )
        .unwrap();

        assert!(!record.envelope.auth_tag.is_empty());
    }

    #[test]
    fn test_different_passwords_different_records() {
        let oprf_seed = generate_oprf_seed();
        let (server_pk, _, _) = create_server_keys();
        let credential_id = b"user@example.com";

        let (client_state1, request1) = client_registration_start(b"password1");
        let response1 =
            server_registration_respond(&oprf_seed, &request1, &server_pk, credential_id).unwrap();
        let record1 = client_registration_finish(client_state1, &response1, None, None).unwrap();

        let (client_state2, request2) = client_registration_start(b"password2");
        let response2 =
            server_registration_respond(&oprf_seed, &request2, &server_pk, credential_id).unwrap();
        let record2 = client_registration_finish(client_state2, &response2, None, None).unwrap();

        assert_ne!(record1.envelope.auth_tag, record2.envelope.auth_tag);
    }
}