//! Hybrid OPRF using Kyber KEM + HMAC-SHA512
//!
//! Since pure lattice-based OPRFs don't have practical implementations,
//! we use a hybrid construction:
//!
//! 1. Client generates ephemeral Kyber keypair and hashes password
//! 2. Client sends blinded_element = (password_hash, ephemeral_pk)
//! 3. Server encapsulates to ephemeral_pk, computes PRF with its secret
//! 4. Server returns (ciphertext, encrypted_prf_output)
//! 5. Client decapsulates, decrypts, derives randomized password

use sha2::{Digest, Sha512};
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::ake::{
    KyberCiphertext, KyberPublicKey, KyberSecretKey, decapsulate, encapsulate, generate_kem_keypair,
};
use crate::error::{OpaqueError, Result};
use crate::kdf::{HASH_LEN, hkdf_expand_fixed};
use crate::mac;
use crate::types::{KYBER_CT_LEN, KYBER_PK_LEN, OprfSeed};

const OPRF_CONTEXT: &[u8] = b"OPAQUE-Lattice-OPRF-v1";
const BLIND_LABEL: &[u8] = b"OPRF-Blind";
const FINALIZE_LABEL: &[u8] = b"OPRF-Finalize";

#[derive(Clone)]
pub struct BlindedElement {
    pub password_hash: [u8; HASH_LEN],
    pub ephemeral_pk: Vec<u8>,
}

impl BlindedElement {
    #[must_use]
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(HASH_LEN + KYBER_PK_LEN);
        bytes.extend_from_slice(&self.password_hash);
        bytes.extend_from_slice(&self.ephemeral_pk);
        bytes
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != HASH_LEN + KYBER_PK_LEN {
            return Err(OpaqueError::InvalidOprfInput);
        }

        let mut password_hash = [0u8; HASH_LEN];
        password_hash.copy_from_slice(&bytes[..HASH_LEN]);
        let ephemeral_pk = bytes[HASH_LEN..].to_vec();

        Ok(Self {
            password_hash,
            ephemeral_pk,
        })
    }
}

#[derive(Clone)]
pub struct EvaluatedElement {
    pub ciphertext: Vec<u8>,
    pub encrypted_output: [u8; HASH_LEN],
}

impl EvaluatedElement {
    #[must_use]
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(KYBER_CT_LEN + HASH_LEN);
        bytes.extend_from_slice(&self.ciphertext);
        bytes.extend_from_slice(&self.encrypted_output);
        bytes
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != KYBER_CT_LEN + HASH_LEN {
            return Err(OpaqueError::InvalidOprfOutput);
        }

        let ciphertext = bytes[..KYBER_CT_LEN].to_vec();
        let mut encrypted_output = [0u8; HASH_LEN];
        encrypted_output.copy_from_slice(&bytes[KYBER_CT_LEN..]);

        Ok(Self {
            ciphertext,
            encrypted_output,
        })
    }
}

#[derive(Zeroize, ZeroizeOnDrop)]
pub struct OprfClient {
    password_hash: [u8; HASH_LEN],
    ephemeral_sk: Vec<u8>,
}

impl OprfClient {
    pub fn blind(password: &[u8]) -> (Self, BlindedElement) {
        let password_hash: [u8; HASH_LEN] = Sha512::digest(password).into();

        let (ephemeral_pk, ephemeral_sk) = generate_kem_keypair();

        let blinded = BlindedElement {
            password_hash,
            ephemeral_pk: ephemeral_pk.as_bytes(),
        };

        let client = Self {
            password_hash,
            ephemeral_sk: ephemeral_sk.as_bytes(),
        };

        (client, blinded)
    }

    pub fn finalize(self, evaluated: &EvaluatedElement) -> Result<[u8; HASH_LEN]> {
        let sk = KyberSecretKey::from_bytes(&self.ephemeral_sk)?;
        let ct = KyberCiphertext::from_bytes(&evaluated.ciphertext)?;

        let shared_secret = decapsulate(&ct, &sk)?;

        let decryption_key: [u8; HASH_LEN] =
            hkdf_expand_fixed(Some(OPRF_CONTEXT), shared_secret.as_bytes(), BLIND_LABEL)?;

        let mut prf_output = [0u8; HASH_LEN];
        for i in 0..HASH_LEN {
            prf_output[i] = evaluated.encrypted_output[i] ^ decryption_key[i];
        }

        let mut finalize_input = Vec::with_capacity(HASH_LEN * 2);
        finalize_input.extend_from_slice(&self.password_hash);
        finalize_input.extend_from_slice(&prf_output);

        let result: [u8; HASH_LEN] =
            hkdf_expand_fixed(Some(OPRF_CONTEXT), &finalize_input, FINALIZE_LABEL)?;

        Ok(result)
    }
}

pub struct OprfServer {
    seed: OprfSeed,
}

impl OprfServer {
    #[must_use]
    pub fn new(seed: OprfSeed) -> Self {
        Self { seed }
    }

    pub fn evaluate(&self, blinded: &BlindedElement) -> Result<EvaluatedElement> {
        let ephemeral_pk = KyberPublicKey::from_bytes(&blinded.ephemeral_pk)?;

        let (shared_secret, ciphertext) = encapsulate(&ephemeral_pk)?;

        let prf_output = mac::compute(self.seed.as_bytes(), &blinded.password_hash);

        let encryption_key: [u8; HASH_LEN] =
            hkdf_expand_fixed(Some(OPRF_CONTEXT), shared_secret.as_bytes(), BLIND_LABEL)?;

        let mut encrypted_output = [0u8; HASH_LEN];
        for i in 0..HASH_LEN {
            encrypted_output[i] = prf_output[i] ^ encryption_key[i];
        }

        Ok(EvaluatedElement {
            ciphertext: ciphertext.as_bytes(),
            encrypted_output,
        })
    }

    pub fn evaluate_with_credential_id(
        &self,
        blinded: &BlindedElement,
        credential_id: &[u8],
    ) -> Result<EvaluatedElement> {
        let ephemeral_pk = KyberPublicKey::from_bytes(&blinded.ephemeral_pk)?;

        let (shared_secret, ciphertext) = encapsulate(&ephemeral_pk)?;

        let mut prf_input = Vec::with_capacity(blinded.password_hash.len() + credential_id.len());
        prf_input.extend_from_slice(&blinded.password_hash);
        prf_input.extend_from_slice(credential_id);
        let prf_output = mac::compute(self.seed.as_bytes(), &prf_input);

        let encryption_key: [u8; HASH_LEN] =
            hkdf_expand_fixed(Some(OPRF_CONTEXT), shared_secret.as_bytes(), BLIND_LABEL)?;

        let mut encrypted_output = [0u8; HASH_LEN];
        for i in 0..HASH_LEN {
            encrypted_output[i] = prf_output[i] ^ encryption_key[i];
        }

        Ok(EvaluatedElement {
            ciphertext: ciphertext.as_bytes(),
            encrypted_output,
        })
    }
}

pub fn server_evaluate(seed: &OprfSeed, blinded: &BlindedElement) -> Result<EvaluatedElement> {
    let server = OprfServer::new(seed.clone());
    server.evaluate(blinded)
}

pub fn client_finalize(client: OprfClient, evaluated: &EvaluatedElement) -> Result<[u8; HASH_LEN]> {
    client.finalize(evaluated)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn random_seed() -> OprfSeed {
        use crate::types::OPRF_SEED_LEN;
        use rand::RngCore;
        let mut bytes = [0u8; OPRF_SEED_LEN];
        rand::thread_rng().fill_bytes(&mut bytes);
        OprfSeed::new(bytes)
    }

    #[test]
    fn test_oprf_roundtrip() {
        let seed = random_seed();
        let password = b"correct horse battery staple";

        let (client, blinded) = OprfClient::blind(password);

        let evaluated = server_evaluate(&seed, &blinded).unwrap();

        let output = client_finalize(client, &evaluated).unwrap();

        assert_eq!(output.len(), HASH_LEN);
    }

    #[test]
    fn test_oprf_deterministic() {
        let seed = random_seed();
        let password = b"test password";

        let (client1, blinded1) = OprfClient::blind(password);
        let evaluated1 = server_evaluate(&seed, &blinded1).unwrap();
        let output1 = client_finalize(client1, &evaluated1).unwrap();

        let (client2, blinded2) = OprfClient::blind(password);
        let evaluated2 = server_evaluate(&seed, &blinded2).unwrap();
        let output2 = client_finalize(client2, &evaluated2).unwrap();

        assert_eq!(output1, output2);
    }

    #[test]
    fn test_oprf_different_passwords() {
        let seed = random_seed();

        let (client1, blinded1) = OprfClient::blind(b"password1");
        let evaluated1 = server_evaluate(&seed, &blinded1).unwrap();
        let output1 = client_finalize(client1, &evaluated1).unwrap();

        let (client2, blinded2) = OprfClient::blind(b"password2");
        let evaluated2 = server_evaluate(&seed, &blinded2).unwrap();
        let output2 = client_finalize(client2, &evaluated2).unwrap();

        assert_ne!(output1, output2);
    }

    #[test]
    fn test_oprf_different_seeds() {
        let seed1 = random_seed();
        let seed2 = random_seed();
        let password = b"same password";

        let (client1, blinded1) = OprfClient::blind(password);
        let evaluated1 = server_evaluate(&seed1, &blinded1).unwrap();
        let output1 = client_finalize(client1, &evaluated1).unwrap();

        let (client2, blinded2) = OprfClient::blind(password);
        let evaluated2 = server_evaluate(&seed2, &blinded2).unwrap();
        let output2 = client_finalize(client2, &evaluated2).unwrap();

        assert_ne!(output1, output2);
    }

    #[test]
    fn test_blinded_element_serialization() {
        let (_, blinded) = OprfClient::blind(b"password");

        let bytes = blinded.to_bytes();
        let restored = BlindedElement::from_bytes(&bytes).unwrap();

        assert_eq!(blinded.password_hash, restored.password_hash);
        assert_eq!(blinded.ephemeral_pk, restored.ephemeral_pk);
    }

    #[test]
    fn test_evaluated_element_serialization() {
        let seed = random_seed();
        let (_, blinded) = OprfClient::blind(b"password");
        let evaluated = server_evaluate(&seed, &blinded).unwrap();

        let bytes = evaluated.to_bytes();
        let restored = EvaluatedElement::from_bytes(&bytes).unwrap();

        assert_eq!(evaluated.ciphertext, restored.ciphertext);
        assert_eq!(evaluated.encrypted_output, restored.encrypted_output);
    }

    #[test]
    fn test_evaluate_with_credential_id() {
        let seed = random_seed();
        let password = b"password";
        let cred_id1 = b"user@example.com";
        let cred_id2 = b"other@example.com";

        let server = OprfServer::new(seed);

        let (client1, blinded1) = OprfClient::blind(password);
        let evaluated1 = server
            .evaluate_with_credential_id(&blinded1, cred_id1)
            .unwrap();
        let output1 = client_finalize(client1, &evaluated1).unwrap();

        let (client2, blinded2) = OprfClient::blind(password);
        let evaluated2 = server
            .evaluate_with_credential_id(&blinded2, cred_id2)
            .unwrap();
        let output2 = client_finalize(client2, &evaluated2).unwrap();

        assert_ne!(output1, output2);
    }
}