use anyhow::{Context, Result};
use argon2::{Argon2, Params};
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Nonce,
};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use totp_rs::{Algorithm, Secret, TOTP};
use zeroize::Zeroize;

const ARGON2_M_COST: u32 = 65536; // 64 MiB
const ARGON2_T_COST: u32 = 3;
const ARGON2_P_COST: u32 = 4;
const ARGON2_OUTPUT_LEN: usize = 32;

const BACKUP_CODE_COUNT: usize = 8;
const BACKUP_CODE_LEN: usize = 8; // 8 alphanumeric chars

/// Encrypted TOTP data stored in user.json. The secret never touches disk in plaintext.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TotpData {
    /// Argon2id salt for KEK derivation (base64)
    pub kek_salt: String,
    /// Nonce for MEK encryption (base64)
    pub mek_nonce: String,
    /// MEK encrypted under KEK via ChaCha20-Poly1305 (base64)
    pub encrypted_mek: String,
    /// Nonce for TOTP secret encryption (base64)
    pub secret_nonce: String,
    /// TOTP secret encrypted under MEK via ChaCha20-Poly1305 (base64)
    pub encrypted_secret: String,
    /// Hashed backup codes (bcrypt), one-time use
    pub backup_codes: Vec<String>,
    /// Recently used TOTP time steps for replay protection
    #[serde(default)]
    pub used_steps: Vec<i64>,
}

/// Result of setup: encrypted data for storage + plaintext for display (shown once).
pub struct SetupResult {
    pub totp_data: TotpData,
    pub secret_base32: String,
    pub qr_svg: String,
    pub backup_codes: Vec<String>,
}

/// Generate a TOTP setup. Returns encrypted data + display values.
pub fn setup(password: &str) -> Result<SetupResult> {
    // Generate the raw TOTP secret (20 bytes = 160 bits, standard for SHA1)
    let mut totp_secret = vec![0u8; 20];
    rand::rngs::OsRng.fill_bytes(&mut totp_secret);

    // Generate MEK (Master Encryption Key)
    let mut mek = [0u8; 32];
    rand::rngs::OsRng.fill_bytes(&mut mek);

    // Derive KEK from password via Argon2id
    let mut kek_salt = [0u8; 16];
    rand::rngs::OsRng.fill_bytes(&mut kek_salt);
    let mut kek = derive_kek(password, &kek_salt)?;

    // Encrypt MEK under KEK
    let mut mek_nonce_bytes = [0u8; 12];
    rand::rngs::OsRng.fill_bytes(&mut mek_nonce_bytes);
    let encrypted_mek = encrypt_chacha(&kek, &mek_nonce_bytes, &mek)?;

    // Encrypt TOTP secret under MEK
    let mut secret_nonce_bytes = [0u8; 12];
    rand::rngs::OsRng.fill_bytes(&mut secret_nonce_bytes);
    let encrypted_secret = encrypt_chacha(&mek, &secret_nonce_bytes, &totp_secret)?;

    // Generate backup codes
    let (plaintext_codes, hashed_codes) = generate_backup_codes()?;

    // Encode the base32 secret for the authenticator app
    let secret_base32 = data_encoding::BASE32_NOPAD.encode(&totp_secret);

    // Generate QR code SVG
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1, // skew
        30,
        Secret::Raw(totp_secret.clone())
            .to_bytes()
            .map_err(|_| anyhow::anyhow!("Invalid TOTP secret"))?,
        Some("Archipelago".to_string()),
        "node".to_string(),
    )
    .context("Failed to create TOTP")?;
    let otpauth_url = totp.get_url();
    let qr_svg = generate_qr_svg(&otpauth_url)?;

    let totp_data = TotpData {
        kek_salt: base64::Engine::encode(&base64::engine::general_purpose::STANDARD, kek_salt),
        mek_nonce: base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            mek_nonce_bytes,
        ),
        encrypted_mek: base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            &encrypted_mek,
        ),
        secret_nonce: base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            secret_nonce_bytes,
        ),
        encrypted_secret: base64::Engine::encode(
            &base64::engine::general_purpose::STANDARD,
            &encrypted_secret,
        ),
        backup_codes: hashed_codes,
        used_steps: Vec::new(),
    };

    // Zeroize sensitive material
    mek.zeroize();
    kek.zeroize();
    totp_secret.zeroize();

    Ok(SetupResult {
        totp_data,
        secret_base32,
        qr_svg,
        backup_codes: plaintext_codes,
    })
}

/// Decrypt the TOTP secret from stored data using the user's password.
/// Returns the raw secret bytes.
pub fn decrypt_secret(data: &TotpData, password: &str) -> Result<Vec<u8>> {
    use base64::Engine;
    let b64 = &base64::engine::general_purpose::STANDARD;

    let kek_salt = b64
        .decode(&data.kek_salt)
        .context("Invalid kek_salt base64")?;
    let mek_nonce = b64
        .decode(&data.mek_nonce)
        .context("Invalid mek_nonce base64")?;
    let encrypted_mek = b64
        .decode(&data.encrypted_mek)
        .context("Invalid encrypted_mek base64")?;
    let secret_nonce = b64
        .decode(&data.secret_nonce)
        .context("Invalid secret_nonce base64")?;
    let encrypted_secret = b64
        .decode(&data.encrypted_secret)
        .context("Invalid encrypted_secret base64")?;

    // Derive KEK
    let mut kek = derive_kek(password, &kek_salt)?;

    // Decrypt MEK
    let mut mek = decrypt_chacha(&kek, &mek_nonce, &encrypted_mek)
        .context("Failed to decrypt MEK — wrong password or corrupt data")?;

    // Decrypt TOTP secret
    let secret = decrypt_chacha(&mek, &secret_nonce, &encrypted_secret)
        .context("Failed to decrypt TOTP secret")?;

    kek.zeroize();
    mek.zeroize();

    Ok(secret)
}

/// Verify a TOTP code against the decrypted secret. Checks ±1 time step window.
pub fn verify_code(secret: &[u8], code: &str, used_steps: &[i64]) -> Result<Option<i64>> {
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        secret.to_vec(),
        None,
        String::new(),
    )
    .context("Failed to create TOTP verifier")?;

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .context("System time error")?
        .as_secs();

    // Check current step and ±1
    for offset in [-1i64, 0, 1] {
        let time = (now as i64) + (offset * 30);
        if time < 0 {
            continue;
        }
        let step = time / 30;
        let expected = totp.generate(time as u64);

        // Constant-time comparison
        if constant_time_eq(code.as_bytes(), expected.as_bytes()) {
            // Replay protection
            if used_steps.contains(&step) {
                return Ok(None); // Code already used
            }
            return Ok(Some(step));
        }
    }

    Ok(None) // No match
}

/// Verify a backup code against the stored bcrypt hashes. Returns the index if valid.
pub fn verify_backup_code(hashed_codes: &[String], code: &str) -> Result<Option<usize>> {
    let normalized = code.replace('-', "").to_uppercase();
    for (i, hash) in hashed_codes.iter().enumerate() {
        if bcrypt::verify(&normalized, hash)? {
            return Ok(Some(i));
        }
    }
    Ok(None)
}

/// Re-encrypt the MEK under a new password (for password change).
pub fn rekey(data: &TotpData, old_password: &str, new_password: &str) -> Result<TotpData> {
    use base64::Engine;
    let b64 = &base64::engine::general_purpose::STANDARD;

    // Decrypt MEK with old password
    let old_kek_salt = b64.decode(&data.kek_salt)?;
    let old_mek_nonce = b64.decode(&data.mek_nonce)?;
    let old_encrypted_mek = b64.decode(&data.encrypted_mek)?;

    let mut old_kek = derive_kek(old_password, &old_kek_salt)?;
    let mut mek = decrypt_chacha(&old_kek, &old_mek_nonce, &old_encrypted_mek)
        .context("Failed to decrypt MEK with old password")?;

    // Re-encrypt MEK with new password
    let mut new_kek_salt = [0u8; 16];
    rand::rngs::OsRng.fill_bytes(&mut new_kek_salt);
    let mut new_kek = derive_kek(new_password, &new_kek_salt)?;

    let mut new_mek_nonce = [0u8; 12];
    rand::rngs::OsRng.fill_bytes(&mut new_mek_nonce);
    let new_encrypted_mek = encrypt_chacha(&new_kek, &new_mek_nonce, &mek)?;

    old_kek.zeroize();
    new_kek.zeroize();
    mek.zeroize();

    Ok(TotpData {
        kek_salt: b64.encode(new_kek_salt),
        mek_nonce: b64.encode(new_mek_nonce),
        encrypted_mek: b64.encode(&new_encrypted_mek),
        // TOTP secret ciphertext unchanged
        secret_nonce: data.secret_nonce.clone(),
        encrypted_secret: data.encrypted_secret.clone(),
        backup_codes: data.backup_codes.clone(),
        used_steps: data.used_steps.clone(),
    })
}

// --- Internal helpers ---

fn derive_kek(password: &str, salt: &[u8]) -> Result<[u8; 32]> {
    let params = Params::new(
        ARGON2_M_COST,
        ARGON2_T_COST,
        ARGON2_P_COST,
        Some(ARGON2_OUTPUT_LEN),
    )
    .map_err(|e| anyhow::anyhow!("Invalid Argon2 params: {}", e))?;
    let argon2 = Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);

    let mut kek = [0u8; 32];
    argon2
        .hash_password_into(password.as_bytes(), salt, &mut kek)
        .map_err(|e| anyhow::anyhow!("Argon2 key derivation failed: {}", e))?;
    Ok(kek)
}

fn encrypt_chacha(key: &[u8; 32], nonce: &[u8], plaintext: &[u8]) -> Result<Vec<u8>> {
    let cipher = ChaCha20Poly1305::new(key.into());
    let nonce = Nonce::from_slice(nonce);
    cipher
        .encrypt(nonce, plaintext)
        .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))
}

fn decrypt_chacha(key: &[u8], nonce: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>> {
    let key: &[u8; 32] = key
        .try_into()
        .map_err(|_| anyhow::anyhow!("Invalid key length"))?;
    let cipher = ChaCha20Poly1305::new(key.into());
    let nonce = Nonce::from_slice(nonce);
    cipher
        .decrypt(nonce, ciphertext)
        .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))
}

fn generate_backup_codes() -> Result<(Vec<String>, Vec<String>)> {
    let charset: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789"; // No 0/O/1/I ambiguity
    let mut plaintext = Vec::with_capacity(BACKUP_CODE_COUNT);
    let mut hashed = Vec::with_capacity(BACKUP_CODE_COUNT);

    for _ in 0..BACKUP_CODE_COUNT {
        let mut code = String::with_capacity(BACKUP_CODE_LEN);
        for _ in 0..BACKUP_CODE_LEN {
            let idx = (rand::random::<u8>() as usize) % charset.len();
            code.push(charset[idx] as char);
        }
        let formatted = format!("{}-{}", &code[..4], &code[4..]);
        let hash = bcrypt::hash(&code, 10)?;
        plaintext.push(formatted);
        hashed.push(hash);
    }

    Ok((plaintext, hashed))
}

fn generate_qr_svg(data: &str) -> Result<String> {
    use qrcode::QrCode;
    let code = QrCode::new(data.as_bytes()).context("Failed to generate QR code")?;
    let svg = code
        .render::<qrcode::render::svg::Color>()
        .min_dimensions(200, 200)
        .quiet_zone(true)
        .build();
    Ok(svg)
}

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)]
mod tests {
    use super::*;

    #[test]
    fn test_setup_and_verify() {
        let password = "TestPassword123!";
        let result = setup(password).unwrap();

        // Decrypt and verify a code
        let secret = decrypt_secret(&result.totp_data, password).unwrap();
        let totp = TOTP::new(
            Algorithm::SHA1,
            6,
            1,
            30,
            secret.clone(),
            None,
            String::new(),
        )
        .unwrap();
        let code = totp.generate_current().unwrap();

        let step = verify_code(&secret, &code, &[]).unwrap();
        assert!(step.is_some(), "Valid TOTP code should verify");

        // Replay: same step should be rejected
        let used_step = step.unwrap();
        let step2 = verify_code(&secret, &code, &[used_step]).unwrap();
        assert!(step2.is_none(), "Replayed code should be rejected");
    }

    #[test]
    fn test_wrong_password_fails() {
        let result = setup("CorrectPassword1!").unwrap();
        let err = decrypt_secret(&result.totp_data, "WrongPassword1!");
        assert!(err.is_err(), "Wrong password should fail decryption");
    }

    #[test]
    fn test_rekey() {
        let old_pw = "OldPassword123!";
        let new_pw = "NewPassword456!";
        let result = setup(old_pw).unwrap();

        // Get original secret
        let original_secret = decrypt_secret(&result.totp_data, old_pw).unwrap();

        // Rekey
        let rekeyed = rekey(&result.totp_data, old_pw, new_pw).unwrap();

        // Old password should fail
        assert!(decrypt_secret(&rekeyed, old_pw).is_err());

        // New password should produce same secret
        let new_secret = decrypt_secret(&rekeyed, new_pw).unwrap();
        assert_eq!(original_secret, new_secret);
    }

    #[test]
    fn test_backup_codes() {
        let result = setup("TestPassword123!").unwrap();
        assert_eq!(result.backup_codes.len(), BACKUP_CODE_COUNT);

        // Verify a backup code works
        let code = &result.backup_codes[0];
        let idx = verify_backup_code(&result.totp_data.backup_codes, code).unwrap();
        assert_eq!(idx, Some(0));

        // Invalid code should fail
        let bad = verify_backup_code(&result.totp_data.backup_codes, "ZZZZ-ZZZZ").unwrap();
        assert!(bad.is_none());
    }
}