use ed25519_dalek::{Signature, Signer, Verifier};
use serde::{Deserialize, Serialize};
use x25519_dalek::{PublicKey, StaticSecret};

use crate::errors::ProtocolError;
use crate::identity::IdentityKeyPair;

/// A signed pre-key (medium-term, rotated periodically).
#[derive(Clone, Serialize, Deserialize)]
pub struct SignedPreKey {
    pub id: u32,
    pub public_key: [u8; 32],
    pub signature: Vec<u8>,
    pub timestamp: i64,
}

impl SignedPreKey {
    /// Verify the signature against the identity signing key.
    pub fn verify(&self, identity_key: &ed25519_dalek::VerifyingKey) -> Result<(), ProtocolError> {
        let sig =
            Signature::from_slice(&self.signature).map_err(|_| ProtocolError::InvalidSignature)?;
        identity_key
            .verify(&self.public_key, &sig)
            .map_err(|_| ProtocolError::PreKeySignatureInvalid)
    }
}

/// A one-time pre-key (used once, then discarded).
pub struct OneTimePreKey {
    pub id: u32,
    pub secret: StaticSecret,
    pub public: PublicKey,
}

/// The public portion of a one-time pre-key (sent to server).
#[derive(Clone, Serialize, Deserialize)]
pub struct OneTimePreKeyPublic {
    pub id: u32,
    pub public_key: [u8; 32],
}

/// A full pre-key bundle that the server stores for a user.
/// Fetched by others to initiate X3DH key exchange.
#[derive(Clone, Serialize, Deserialize)]
pub struct PreKeyBundle {
    pub identity_key: [u8; 32], // Ed25519 verifying key bytes
    pub identity_encryption_key: [u8; 32], // X25519 public key bytes
    pub signed_pre_key: SignedPreKey,
    pub one_time_pre_key: Option<OneTimePreKeyPublic>,
}

/// Generate a signed pre-key.
pub fn generate_signed_pre_key(identity: &IdentityKeyPair, id: u32) -> (StaticSecret, SignedPreKey) {
    let secret = StaticSecret::random_from_rng(rand::rngs::OsRng);
    let public = PublicKey::from(&secret);
    let signature = identity.signing.sign(public.as_bytes());

    let spk = SignedPreKey {
        id,
        public_key: *public.as_bytes(),
        signature: signature.to_bytes().to_vec(),
        timestamp: chrono::Utc::now().timestamp(),
    };
    (secret, spk)
}

/// Generate a batch of one-time pre-keys.
pub fn generate_one_time_pre_keys(start_id: u32, count: u32) -> Vec<OneTimePreKey> {
    (start_id..start_id + count)
        .map(|id| {
            let secret = StaticSecret::random_from_rng(rand::rngs::OsRng);
            let public = PublicKey::from(&secret);
            OneTimePreKey {
                id,
                secret,
                public,
            }
        })
        .collect()
}

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

    #[test]
    fn signed_pre_key_verify() {
        let seed = Seed::generate();
        let identity = seed.derive_identity();
        let (_secret, spk) = generate_signed_pre_key(&identity, 1);
        let pub_id = identity.public_identity();
        assert!(spk.verify(&pub_id.signing).is_ok());
    }

    #[test]
    fn signed_pre_key_reject_tampered() {
        let seed = Seed::generate();
        let identity = seed.derive_identity();
        let (_secret, mut spk) = generate_signed_pre_key(&identity, 1);
        spk.public_key[0] ^= 0xff; // tamper
        let pub_id = identity.public_identity();
        assert!(spk.verify(&pub_id.signing).is_err());
    }

    #[test]
    fn generate_otpks() {
        let keys = generate_one_time_pre_keys(0, 10);
        assert_eq!(keys.len(), 10);
        // All public keys should be unique
        let pubs: Vec<_> = keys.iter().map(|k| *k.public.as_bytes()).collect();
        let unique: std::collections::HashSet<_> = pubs.iter().collect();
        assert_eq!(unique.len(), 10);
    }
}