//! Account recovery — recovery code generation and encrypted backup bundles.
//!
//! # Design
//!
//! Recovery codes are 8 alphanumeric strings of 6 characters each (~31 bits
//! entropy per code). Any single code is sufficient to recover the account.
//!
//! A recovery key is derived from each code via Argon2id. The identity seed
//! and conversation metadata are encrypted into a [`RecoveryBundle`] using
//! ChaCha20-Poly1305. The bundle is uploaded to the server, keyed by
//! `SHA-256(recovery_token)` — the server never sees plaintext codes.
//!
//! # Security properties
//!
//! - Recovery codes are shown once and never stored in plaintext.
//! - The server is zero-knowledge — it stores only encrypted blobs.
//! - Code validation uses constant-time comparison.
//! - All key material is zeroized on drop.

use argon2::{Algorithm, Argon2, Params, Version};
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Key, Nonce,
};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use zeroize::Zeroizing;

use crate::error::CoreError;

/// Number of recovery codes generated per setup.
pub const RECOVERY_CODE_COUNT: usize = 8;

/// Length of each recovery code (alphanumeric characters).
const CODE_LENGTH: usize = 6;

/// Maximum bundle size (64 KiB).
pub const MAX_BUNDLE_SIZE: usize = 64 * 1024;

/// Argon2id parameters for recovery key derivation.
const ARGON2_M_COST: u32 = 19 * 1024; // 19 MiB
const ARGON2_T_COST: u32 = 2;
const ARGON2_P_COST: u32 = 1;

/// Alphanumeric character set for recovery codes (uppercase + digits, no
/// ambiguous characters 0/O, 1/I/L).
const CODE_ALPHABET: &[u8] = b"23456789ABCDEFGHJKMNPQRSTUVWXYZ";

/// An encrypted recovery bundle stored on the server.
///
/// The server stores this keyed by `token_hash` (SHA-256 of a recovery token
/// derived from the code). The server cannot decrypt it.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecoveryBundle {
    /// SHA-256 of the recovery token (used as server-side lookup key).
    pub token_hash: Vec<u8>,
    /// Random 16-byte salt for Argon2id key derivation.
    pub salt: Vec<u8>,
    /// Random 12-byte nonce for ChaCha20-Poly1305.
    pub nonce: Vec<u8>,
    /// Encrypted payload: bincode-serialised `RecoveryPayload`.
    pub ciphertext: Vec<u8>,
}

/// The plaintext payload inside a recovery bundle.
#[derive(Debug, Serialize, Deserialize)]
pub struct RecoveryPayload {
    /// Ed25519 identity seed (32 bytes).
    pub identity_seed: [u8; 32],
    /// List of conversation/group IDs the user was part of (for rejoin).
    pub conversation_ids: Vec<Vec<u8>>,
}

/// Result of recovery code generation.
pub struct RecoverySetup {
    /// The 8 recovery codes to show to the user (shown once, never stored).
    pub codes: Vec<String>,
    /// Encrypted bundles — one per code — to upload to the server.
    pub bundles: Vec<RecoveryBundle>,
}

/// Generate a single random recovery code.
fn generate_code(rng: &mut impl RngCore) -> String {
    let mut code = String::with_capacity(CODE_LENGTH);
    for _ in 0..CODE_LENGTH {
        let idx = (rng.next_u32() as usize) % CODE_ALPHABET.len();
        code.push(CODE_ALPHABET[idx] as char);
    }
    code
}

/// Derive a 32-byte recovery token from a code (used for server-side lookup).
/// The token is `SHA-256("qpq-recovery-token:" || code)`.
fn derive_recovery_token(code: &str) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(b"qpq-recovery-token:");
    hasher.update(code.as_bytes());
    hasher.finalize().into()
}

/// Derive a 32-byte encryption key from a code and salt via Argon2id.
fn derive_recovery_key(code: &str, salt: &[u8]) -> Result<Zeroizing<[u8; 32]>, CoreError> {
    let params = Params::new(ARGON2_M_COST, ARGON2_T_COST, ARGON2_P_COST, Some(32))
        .map_err(|e| CoreError::Io(format!("argon2 params: {e}")))?;
    let argon2 = Argon2::new(Algorithm::Argon2id, Version::default(), params);
    let mut key = Zeroizing::new([0u8; 32]);
    argon2
        .hash_password_into(code.as_bytes(), salt, &mut *key)
        .map_err(|e| CoreError::Io(format!("argon2 recovery key derivation: {e}")))?;
    Ok(key)
}

/// Generate recovery codes and encrypted bundles for an identity.
///
/// Returns a `RecoverySetup` containing:
/// - `codes`: 8 recovery codes to display to the user (once).
/// - `bundles`: 8 encrypted recovery bundles (one per code) to upload to the server.
///
/// Each code independently decrypts its corresponding bundle.
pub fn generate_recovery_codes(
    identity_seed: &[u8; 32],
    conversation_ids: &[Vec<u8>],
) -> Result<RecoverySetup, CoreError> {
    let mut rng = rand::rngs::OsRng;

    let payload = RecoveryPayload {
        identity_seed: *identity_seed,
        conversation_ids: conversation_ids.to_vec(),
    };
    let plaintext = bincode::serialize(&payload)
        .map_err(|e| CoreError::Io(format!("serialize recovery payload: {e}")))?;

    let mut codes = Vec::with_capacity(RECOVERY_CODE_COUNT);
    let mut bundles = Vec::with_capacity(RECOVERY_CODE_COUNT);

    for _ in 0..RECOVERY_CODE_COUNT {
        let code = generate_code(&mut rng);

        // Derive the server-side lookup token.
        let token = derive_recovery_token(&code);
        let token_hash = Sha256::digest(token).to_vec();

        // Derive encryption key from code.
        let mut salt = [0u8; 16];
        rng.fill_bytes(&mut salt);

        let key = derive_recovery_key(&code, &salt)?;
        let cipher = ChaCha20Poly1305::new(Key::from_slice(&*key));

        let mut nonce_bytes = [0u8; 12];
        rng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let ciphertext = cipher
            .encrypt(nonce, plaintext.as_slice())
            .map_err(|e| CoreError::Io(format!("recovery bundle encryption: {e}")))?;

        bundles.push(RecoveryBundle {
            token_hash,
            salt: salt.to_vec(),
            nonce: nonce_bytes.to_vec(),
            ciphertext,
        });
        codes.push(code);
    }

    Ok(RecoverySetup { codes, bundles })
}

/// Recover an identity seed from a recovery code and encrypted bundle.
///
/// Returns the decrypted `RecoveryPayload` on success.
pub fn recover_from_bundle(
    code: &str,
    bundle: &RecoveryBundle,
) -> Result<RecoveryPayload, CoreError> {
    // Validate bundle structure.
    if bundle.salt.len() != 16 {
        return Err(CoreError::Io(format!(
            "invalid recovery bundle salt length: {}",
            bundle.salt.len()
        )));
    }
    if bundle.nonce.len() != 12 {
        return Err(CoreError::Io(format!(
            "invalid recovery bundle nonce length: {}",
            bundle.nonce.len()
        )));
    }

    // Derive encryption key from code.
    let key = derive_recovery_key(code, &bundle.salt)?;
    let cipher = ChaCha20Poly1305::new(Key::from_slice(&*key));
    let nonce = Nonce::from_slice(&bundle.nonce);

    let plaintext = cipher
        .decrypt(nonce, bundle.ciphertext.as_slice())
        .map_err(|_| CoreError::Io("recovery bundle decryption failed (wrong code?)".into()))?;

    let payload: RecoveryPayload = bincode::deserialize(&plaintext)
        .map_err(|e| CoreError::Io(format!("deserialize recovery payload: {e}")))?;

    Ok(payload)
}

/// Compute the token hash for a recovery code (for server-side lookup).
///
/// This is `SHA-256(SHA-256("qpq-recovery-token:" || code))`.
pub fn recovery_token_hash(code: &str) -> Vec<u8> {
    let token = derive_recovery_token(code);
    Sha256::digest(token).to_vec()
}

/// Constant-time comparison of two byte slices.
///
/// Returns `true` if the slices are equal, using constant-time comparison
/// to prevent timing side-channels on recovery code validation.
pub fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff = 0u8;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}

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

    #[test]
    fn generate_codes_produces_correct_count() {
        let seed = [42u8; 32];
        let setup = generate_recovery_codes(&seed, &[]).unwrap();
        assert_eq!(setup.codes.len(), RECOVERY_CODE_COUNT);
        assert_eq!(setup.bundles.len(), RECOVERY_CODE_COUNT);
    }

    #[test]
    fn codes_are_correct_length_and_alphabet() {
        let seed = [7u8; 32];
        let setup = generate_recovery_codes(&seed, &[]).unwrap();
        for code in &setup.codes {
            assert_eq!(code.len(), CODE_LENGTH);
            for ch in code.chars() {
                assert!(
                    CODE_ALPHABET.contains(&(ch as u8)),
                    "invalid char '{ch}' in code"
                );
            }
        }
    }

    #[test]
    fn codes_are_unique() {
        let seed = [1u8; 32];
        let setup = generate_recovery_codes(&seed, &[]).unwrap();
        let mut seen = std::collections::HashSet::new();
        for code in &setup.codes {
            assert!(seen.insert(code.clone()), "duplicate code: {code}");
        }
    }

    #[test]
    fn recover_roundtrip() {
        let seed = [99u8; 32];
        let conv_ids = vec![vec![1, 2, 3], vec![4, 5, 6]];
        let setup = generate_recovery_codes(&seed, &conv_ids).unwrap();

        // Each code should decrypt its corresponding bundle.
        for (i, code) in setup.codes.iter().enumerate() {
            let payload = recover_from_bundle(code, &setup.bundles[i]).unwrap();
            assert_eq!(payload.identity_seed, seed);
            assert_eq!(payload.conversation_ids, conv_ids);
        }
    }

    #[test]
    fn wrong_code_fails() {
        let seed = [50u8; 32];
        let setup = generate_recovery_codes(&seed, &[]).unwrap();
        let result = recover_from_bundle("WRONG1", &setup.bundles[0]);
        assert!(result.is_err());
    }

    #[test]
    fn code_does_not_decrypt_other_bundle() {
        let seed = [88u8; 32];
        let setup = generate_recovery_codes(&seed, &[]).unwrap();
        // Code 0 should NOT decrypt bundle 1 (different salt/nonce/key).
        let result = recover_from_bundle(&setup.codes[0], &setup.bundles[1]);
        assert!(result.is_err());
    }

    #[test]
    fn token_hash_is_deterministic() {
        let hash1 = recovery_token_hash("ABC123");
        let hash2 = recovery_token_hash("ABC123");
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn token_hash_differs_for_different_codes() {
        let hash1 = recovery_token_hash("ABC123");
        let hash2 = recovery_token_hash("XYZ789");
        assert_ne!(hash1, hash2);
    }

    #[test]
    fn constant_time_eq_works() {
        assert!(constant_time_eq(b"hello", b"hello"));
        assert!(!constant_time_eq(b"hello", b"world"));
        assert!(!constant_time_eq(b"hello", b"hell"));
        assert!(constant_time_eq(b"", b""));
    }

    #[test]
    fn invalid_bundle_salt_rejected() {
        let bundle = RecoveryBundle {
            token_hash: vec![0; 32],
            salt: vec![0; 8], // wrong length
            nonce: vec![0; 12],
            ciphertext: vec![0; 32],
        };
        assert!(recover_from_bundle("ABC123", &bundle).is_err());
    }

    #[test]
    fn invalid_bundle_nonce_rejected() {
        let bundle = RecoveryBundle {
            token_hash: vec![0; 32],
            salt: vec![0; 16],
            nonce: vec![0; 8], // wrong length
            ciphertext: vec![0; 32],
        };
        assert!(recover_from_bundle("ABC123", &bundle).is_err());
    }
}