use rand::RngCore;

use crate::ake::{
    DilithiumPublicKey, DilithiumSecretKey, KyberCiphertext, KyberPublicKey, decapsulate,
    encapsulate, generate_kem_keypair, sign, verify,
};
use crate::envelope;
use crate::error::{OpaqueError, Result};
use crate::kdf::{HASH_LEN, hkdf_expand_fixed, labels};
use crate::mac;
use crate::oprf::{BlindedElement, EvaluatedElement, OprfClient, OprfServer};
use crate::types::{
    AuthRequest, AuthResponse, CredentialRequest, CredentialResponse, Envelope, KE1, KE2, KE3,
    NONCE_LEN, OprfSeed, RegistrationRecord, ServerPrivateKey, ServerPublicKey, SessionKey,
};

const HANDSHAKE_CONTEXT: &[u8] = b"OPAQUE-Lattice-Handshake-v1";

pub struct ClientLoginState {
    oprf_client: OprfClient,
    client_nonce: [u8; NONCE_LEN],
    client_kem_pk: Vec<u8>,
    client_kem_sk: Vec<u8>,
}

pub struct ServerLoginState {
    expected_client_mac: [u8; HASH_LEN],
    session_key: [u8; HASH_LEN],
}

pub fn client_login_start(password: &[u8]) -> (ClientLoginState, KE1) {
    let (oprf_client, blinded) = OprfClient::blind(password);

    let mut client_nonce = [0u8; NONCE_LEN];
    rand::rng().fill_bytes(&mut client_nonce);

    let (client_kem_pk, client_kem_sk) = generate_kem_keypair();

    #[cfg(feature = "debug-trace")]
    {
        eprintln!("[LOGIN] client_login_start");
        eprintln!("  client_nonce: {:02x?}", &client_nonce[..8]);
        eprintln!("  client_kem_pk len: {}", client_kem_pk.as_bytes().len());
    }

    let ke1 = KE1 {
        credential_request: CredentialRequest {
            blinded_element: blinded.to_bytes(),
        },
        auth_request: AuthRequest {
            client_nonce,
            client_kem_pk: client_kem_pk.as_bytes(),
        },
    };

    let state = ClientLoginState {
        oprf_client,
        client_nonce,
        client_kem_pk: client_kem_pk.as_bytes(),
        client_kem_sk: client_kem_sk.as_bytes(),
    };

    (state, ke1)
}

pub fn server_login_respond(
    oprf_seed: &OprfSeed,
    server_public_key: &ServerPublicKey,
    server_private_key: &ServerPrivateKey,
    record: &RegistrationRecord,
    credential_id: &[u8],
    ke1: &KE1,
) -> Result<(ServerLoginState, KE2)> {
    #[cfg(feature = "debug-trace")]
    eprintln!("[LOGIN] server_login_respond starting");

    let blinded = BlindedElement::from_bytes(&ke1.credential_request.blinded_element)?;

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

    #[cfg(feature = "debug-trace")]
    eprintln!("  OPRF evaluation complete");

    let mut masking_nonce = [0u8; NONCE_LEN];
    rand::rng().fill_bytes(&mut masking_nonce);

    let envelope_bytes = serialize_envelope(&record.envelope);
    let to_mask = [
        &server_public_key.kem_pk[..],
        &server_public_key.sig_pk[..],
        &envelope_bytes[..],
    ]
    .concat();
    let masked_response = envelope::mask_response(&record.masking_key, &masking_nonce, &to_mask)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  masked_response len: {}", masked_response.len());

    let mut server_nonce = [0u8; NONCE_LEN];
    rand::rng().fill_bytes(&mut server_nonce);

    let (server_kem_pk, _server_kem_sk) = generate_kem_keypair();

    let client_kem_pk = KyberPublicKey::from_bytes(&ke1.auth_request.client_kem_pk)?;
    let (shared_secret, server_kem_ct) = encapsulate(&client_kem_pk)?;

    #[cfg(feature = "debug-trace")]
    {
        eprintln!("  shared_secret: {:02x?}", &shared_secret.as_bytes()[..8]);
        eprintln!("  server_nonce: {:02x?}", &server_nonce[..8]);
    }

    let transcript = build_transcript(
        &ke1.auth_request.client_nonce,
        &ke1.auth_request.client_kem_pk,
        &server_nonce,
        &server_kem_pk.as_bytes(),
        &server_kem_ct.as_bytes(),
    );

    #[cfg(feature = "debug-trace")]
    eprintln!("  transcript len: {}", transcript.len());

    let (session_key, server_mac_key, client_mac_key) =
        derive_keys(shared_secret.as_bytes(), &transcript)?;

    #[cfg(feature = "debug-trace")]
    {
        eprintln!("  session_key: {:02x?}", &session_key[..8]);
        eprintln!("  server_mac_key: {:02x?}", &server_mac_key[..8]);
        eprintln!("  client_mac_key: {:02x?}", &client_mac_key[..8]);
    }

    let server_mac = mac::compute(&server_mac_key, &transcript);
    let expected_client_mac = mac::compute(&client_mac_key, &transcript);

    #[cfg(feature = "debug-trace")]
    eprintln!("  server_mac: {:02x?}", &server_mac[..8]);

    let sig_sk = DilithiumSecretKey::from_bytes(&server_private_key.sig_sk)?;
    let signature_data = build_signature_data(
        &server_nonce,
        &server_kem_pk.as_bytes(),
        &server_kem_ct.as_bytes(),
        &server_mac,
    );
    let server_signature = sign(&signature_data, &sig_sk);

    #[cfg(feature = "debug-trace")]
    eprintln!(
        "  server_signature len: {}",
        server_signature.as_bytes().len()
    );

    let ke2 = KE2 {
        credential_response: CredentialResponse {
            evaluated_element: evaluated.to_bytes(),
            masking_nonce,
            masked_response,
        },
        auth_response: AuthResponse {
            server_nonce,
            server_kem_pk: server_kem_pk.as_bytes(),
            server_kem_ct: server_kem_ct.as_bytes(),
            server_mac: server_mac.to_vec(),
            server_signature: server_signature.as_bytes(),
        },
    };

    let state = ServerLoginState {
        expected_client_mac,
        session_key,
    };

    #[cfg(feature = "debug-trace")]
    eprintln!("[LOGIN] server_login_respond complete");

    Ok((state, ke2))
}

pub fn client_login_finish(
    state: ClientLoginState,
    ke2: &KE2,
    server_identity: Option<&[u8]>,
    client_identity: Option<&[u8]>,
) -> Result<(KE3, SessionKey)> {
    #[cfg(feature = "debug-trace")]
    eprintln!("[LOGIN] client_login_finish starting");

    let evaluated = EvaluatedElement::from_bytes(&ke2.credential_response.evaluated_element)?;
    let randomized_pwd = state.oprf_client.finalize(&evaluated)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  randomized_pwd: {:02x?}", &randomized_pwd[..8]);

    let masking_key = envelope::create_masking_key(&randomized_pwd)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  masking_key: {:02x?}", &masking_key[..8]);

    let unmasked = envelope::unmask_response(
        &masking_key,
        &ke2.credential_response.masking_nonce,
        &ke2.credential_response.masked_response,
    )?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  unmasked len: {}", unmasked.len());

    let (server_public_key, envelope) = parse_unmasked_response(&unmasked)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  envelope nonce: {:02x?}", &envelope.nonce[..8]);

    envelope::recover(
        &randomized_pwd,
        &server_public_key,
        &envelope,
        server_identity,
        client_identity,
    )?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  envelope recovered successfully");

    if !server_public_key.sig_pk.is_empty() {
        let sig_pk = DilithiumPublicKey::from_bytes(&server_public_key.sig_pk)?;
        let signature_data = build_signature_data(
            &ke2.auth_response.server_nonce,
            &ke2.auth_response.server_kem_pk,
            &ke2.auth_response.server_kem_ct,
            &ke2.auth_response.server_mac,
        );
        let signature =
            crate::ake::DilithiumSignature::from_bytes(&ke2.auth_response.server_signature)?;
        verify(&signature_data, &signature, &sig_pk)?;

        #[cfg(feature = "debug-trace")]
        eprintln!("  server signature verified");
    }

    let server_kem_ct = KyberCiphertext::from_bytes(&ke2.auth_response.server_kem_ct)?;
    let client_kem_sk = crate::ake::KyberSecretKey::from_bytes(&state.client_kem_sk)?;
    let shared_secret = decapsulate(&server_kem_ct, &client_kem_sk)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  shared_secret: {:02x?}", &shared_secret.as_bytes()[..8]);

    let transcript = build_transcript(
        &state.client_nonce,
        &state.client_kem_pk,
        &ke2.auth_response.server_nonce,
        &ke2.auth_response.server_kem_pk,
        &ke2.auth_response.server_kem_ct,
    );

    #[cfg(feature = "debug-trace")]
    eprintln!("  transcript len: {}", transcript.len());

    let (session_key, server_mac_key, client_mac_key) =
        derive_keys(shared_secret.as_bytes(), &transcript)?;

    #[cfg(feature = "debug-trace")]
    {
        eprintln!("  session_key: {:02x?}", &session_key[..8]);
        eprintln!("  server_mac_key: {:02x?}", &server_mac_key[..8]);
    }

    mac::verify(&server_mac_key, &transcript, &ke2.auth_response.server_mac)?;

    #[cfg(feature = "debug-trace")]
    eprintln!("  server MAC verified");

    let client_mac = mac::compute(&client_mac_key, &transcript);

    let ke3 = KE3 {
        client_mac: client_mac.to_vec(),
    };

    #[cfg(feature = "debug-trace")]
    eprintln!("[LOGIN] client_login_finish complete");

    Ok((ke3, SessionKey::new(session_key)))
}

pub fn server_login_finish(state: ServerLoginState, ke3: &KE3) -> Result<SessionKey> {
    use subtle::ConstantTimeEq;

    #[cfg(feature = "debug-trace")]
    eprintln!("[LOGIN] server_login_finish");

    if ke3.client_mac.len() != HASH_LEN {
        return Err(OpaqueError::MacVerificationFailed);
    }

    if state.expected_client_mac.ct_eq(&ke3.client_mac).into() {
        #[cfg(feature = "debug-trace")]
        eprintln!("  client MAC verified");
        Ok(SessionKey::new(state.session_key))
    } else {
        Err(OpaqueError::MacVerificationFailed)
    }
}

fn derive_keys(
    shared_secret: &[u8],
    transcript: &[u8],
) -> Result<([u8; HASH_LEN], [u8; HASH_LEN], [u8; HASH_LEN])> {
    let mut ikm = Vec::with_capacity(shared_secret.len() + transcript.len());
    ikm.extend_from_slice(shared_secret);
    ikm.extend_from_slice(transcript);

    let handshake_secret: [u8; HASH_LEN] =
        hkdf_expand_fixed(Some(HANDSHAKE_CONTEXT), &ikm, labels::HANDSHAKE_SECRET)?;

    let session_key: [u8; HASH_LEN] = hkdf_expand_fixed(
        Some(HANDSHAKE_CONTEXT),
        &handshake_secret,
        labels::SESSION_KEY,
    )?;

    let server_mac_key: [u8; HASH_LEN] = hkdf_expand_fixed(
        Some(HANDSHAKE_CONTEXT),
        &handshake_secret,
        labels::SERVER_MAC,
    )?;

    let client_mac_key: [u8; HASH_LEN] = hkdf_expand_fixed(
        Some(HANDSHAKE_CONTEXT),
        &handshake_secret,
        labels::CLIENT_MAC,
    )?;

    Ok((session_key, server_mac_key, client_mac_key))
}

fn build_transcript(
    client_nonce: &[u8],
    client_kem_pk: &[u8],
    server_nonce: &[u8],
    server_kem_pk: &[u8],
    server_kem_ct: &[u8],
) -> Vec<u8> {
    let mut transcript = Vec::new();
    transcript.extend_from_slice(client_nonce);
    transcript.extend_from_slice(client_kem_pk);
    transcript.extend_from_slice(server_nonce);
    transcript.extend_from_slice(server_kem_pk);
    transcript.extend_from_slice(server_kem_ct);
    transcript
}

fn build_signature_data(
    server_nonce: &[u8],
    server_kem_pk: &[u8],
    server_kem_ct: &[u8],
    server_mac: &[u8],
) -> Vec<u8> {
    let mut data = Vec::new();
    data.extend_from_slice(server_nonce);
    data.extend_from_slice(server_kem_pk);
    data.extend_from_slice(server_kem_ct);
    data.extend_from_slice(server_mac);
    data
}

fn serialize_envelope(envelope: &Envelope) -> Vec<u8> {
    let mut bytes = Vec::new();
    bytes.extend_from_slice(&envelope.nonce);
    bytes.extend_from_slice(&envelope.auth_tag);
    bytes
}

fn parse_unmasked_response(data: &[u8]) -> Result<(ServerPublicKey, Envelope)> {
    use crate::types::{DILITHIUM_PK_LEN, KYBER_PK_LEN};

    let min_len = KYBER_PK_LEN + DILITHIUM_PK_LEN + NONCE_LEN + HASH_LEN;
    if data.len() < min_len {
        return Err(OpaqueError::Deserialization(
            "Unmasked response too short".into(),
        ));
    }

    let server_kem_pk = data[..KYBER_PK_LEN].to_vec();
    let server_sig_pk = data[KYBER_PK_LEN..KYBER_PK_LEN + DILITHIUM_PK_LEN].to_vec();
    let remaining = &data[KYBER_PK_LEN + DILITHIUM_PK_LEN..];

    let mut envelope_nonce = [0u8; NONCE_LEN];
    envelope_nonce.copy_from_slice(&remaining[..NONCE_LEN]);
    let auth_tag = remaining[NONCE_LEN..].to_vec();

    let server_public_key = ServerPublicKey::new(server_kem_pk, server_sig_pk);
    let envelope = Envelope::new(envelope_nonce, auth_tag);

    Ok((server_public_key, envelope))
}

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

    fn create_server_keys() -> (ServerPublicKey, ServerPrivateKey) {
        let (kem_pk, kem_sk) = generate_kem_keypair();
        let (sig_pk, sig_sk) = generate_sig_keypair();
        let public_key = ServerPublicKey::new(kem_pk.as_bytes(), sig_pk.as_bytes());
        let private_key = ServerPrivateKey::new(kem_sk.as_bytes(), sig_sk.as_bytes());
        (public_key, private_key)
    }

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

        let (reg_state, reg_request) = client_registration_start(password);
        let reg_response =
            server_registration_respond(&oprf_seed, &reg_request, &server_pk, credential_id)
                .unwrap();
        let record = client_registration_finish(reg_state, &reg_response, None, None).unwrap();

        let (client_state, ke1) = client_login_start(password);

        let (server_state, ke2) = server_login_respond(
            &oprf_seed,
            &server_pk,
            &server_sk,
            &record,
            credential_id,
            &ke1,
        )
        .unwrap();

        assert!(
            !ke2.auth_response.server_signature.is_empty(),
            "Server signature should be present"
        );

        let (ke3, client_session_key) =
            client_login_finish(client_state, &ke2, None, None).unwrap();

        let server_session_key = server_login_finish(server_state, &ke3).unwrap();

        assert_eq!(client_session_key.as_bytes(), server_session_key.as_bytes());
    }

    #[test]
    fn test_wrong_password_fails() {
        let oprf_seed = generate_oprf_seed();
        let (server_pk, server_sk) = create_server_keys();
        let credential_id = b"user@example.com";
        let correct_password = b"correct password";
        let wrong_password = b"wrong password";

        let (reg_state, reg_request) = client_registration_start(correct_password);
        let reg_response =
            server_registration_respond(&oprf_seed, &reg_request, &server_pk, credential_id)
                .unwrap();
        let record = client_registration_finish(reg_state, &reg_response, None, None).unwrap();

        let (client_state, ke1) = client_login_start(wrong_password);

        let (_, ke2) = server_login_respond(
            &oprf_seed,
            &server_pk,
            &server_sk,
            &record,
            credential_id,
            &ke1,
        )
        .unwrap();

        let result = client_login_finish(client_state, &ke2, None, None);
        assert!(result.is_err());
    }

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

        let (reg_state, reg_request) = client_registration_start(password);
        let reg_response =
            server_registration_respond(&oprf_seed, &reg_request, &server_pk, credential_id)
                .unwrap();
        let record = client_registration_finish(reg_state, &reg_response, None, None).unwrap();

        let (client_state, ke1) = client_login_start(password);

        let (_, mut ke2) = server_login_respond(
            &oprf_seed,
            &server_pk,
            &server_sk,
            &record,
            credential_id,
            &ke1,
        )
        .unwrap();

        ke2.auth_response.server_signature[0] ^= 0xFF;

        let result = client_login_finish(client_state, &ke2, None, None);
        assert!(
            result.is_err(),
            "Tampered signature should fail verification"
        );
    }
}