//! SPAKE2+ implementation for TPAP protocol.
//!
//! SPAKE2+ is an augmented Password-Authenticated Key Exchange (PAKE) protocol
//! as defined in RFC 9383. This implementation uses P-256 (secp256r1) as the
//! elliptic curve group.
//!
//! # Protocol Overview
//!
//! SPAKE2+ allows a client (prover) who knows a password to establish a shared
//! secret with a server (verifier) who only stores a password verifier.
//!
//! 1. **Offline Registration**: Password is derived into (w0, w1) using PBKDF2,
//!    and the verifier stores w0 and L = w1*G.
//!
//! 2. **Online Authentication**:
//!    - Prover generates ephemeral x, computes X = x*G + w0*M (shareP)
//!    - Verifier generates ephemeral y, computes Y = y*G + w0*N (shareV)
//!    - Both derive shared secrets Z and V from ECDH
//!    - Key schedule derives confirmation keys and shared secret
//!
//! # TP-Link TPAP Usage
//!
//! TP-Link's TPAP protocol uses SPAKE2+ with some customizations:
//! - Credentials are transformed through various password_shadow algorithms
//! - Uses P-256 with standard M and N generator points from RFC 9383
//! - Cipher suites: 1=SHA256+HMAC, 2=SHA512+HMAC, 8=SHA256+CMAC, etc.

use hmac::{Hmac, Mac};
use p256::elliptic_curve::ops::MulByGenerator;
use p256::elliptic_curve::sec1::{FromEncodedPoint, ToEncodedPoint};
use p256::{AffinePoint, ProjectivePoint, Scalar};
use sha2::{Digest, Sha256, Sha512};

/// Fixed M point for P-256 from RFC 9383.
/// Compressed SEC1 encoding.
const P256_M_COMPRESSED: [u8; 33] =
    hex_literal::hex!("02886e2f97ace46e55ba9dd7242579f2993b64e16ef3dcab95afd497333d8fa12f");

/// Fixed N point for P-256 from RFC 9383.
/// Compressed SEC1 encoding.
const P256_N_COMPRESSED: [u8; 33] =
    hex_literal::hex!("03d8bbd6c639c62937b04d997f38c3770719c629d7014d49a24b4f98baa1292b49");

/// PAKE context tag used in transcript hashing.
const PAKE_CONTEXT_TAG: &[u8] = b"PAKE V1";

/// SPAKE2+ cipher suite configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Spake2PlusCipherSuite {
    /// Suite 1: SHA256 for hashing, HMAC-SHA256 for confirmation
    #[default]
    P256Sha256HmacSha256 = 1,
    /// Suite 2: SHA512 for hashing, HMAC-SHA512 for confirmation
    P256Sha512HmacSha512 = 2,
    /// Suite 8: SHA256 for hashing, CMAC-AES for confirmation
    P256Sha256CmacAes = 8,
    /// Suite 9: SHA512 for hashing, CMAC-AES for confirmation
    P256Sha512CmacAes = 9,
}

impl Spake2PlusCipherSuite {
    /// Parse cipher suite from integer.
    pub fn from_u8(value: u8) -> Option<Self> {
        match value {
            1 => Some(Self::P256Sha256HmacSha256),
            2 => Some(Self::P256Sha512HmacSha512),
            8 => Some(Self::P256Sha256CmacAes),
            9 => Some(Self::P256Sha512CmacAes),
            _ => None,
        }
    }

    /// Returns the hash algorithm name for HKDF.
    pub fn hash_name(&self) -> &'static str {
        match self {
            Self::P256Sha256HmacSha256 | Self::P256Sha256CmacAes => "SHA256",
            Self::P256Sha512HmacSha512 | Self::P256Sha512CmacAes => "SHA512",
        }
    }

    /// Returns the digest length in bytes.
    pub fn digest_len(&self) -> usize {
        match self {
            Self::P256Sha256HmacSha256 | Self::P256Sha256CmacAes => 32,
            Self::P256Sha512HmacSha512 | Self::P256Sha512CmacAes => 64,
        }
    }

    /// Returns whether this suite uses CMAC for confirmation.
    pub fn uses_cmac(&self) -> bool {
        matches!(self, Self::P256Sha256CmacAes | Self::P256Sha512CmacAes)
    }

    /// Returns the MAC length for confirmation.
    pub fn mac_len(&self) -> usize {
        if self.uses_cmac() { 16 } else { 32 }
    }
}

/// SPAKE2+ prover (client) state machine.
pub struct Spake2PlusProver {
    /// Password-derived w0 scalar (mod group order).
    w0: Scalar,
    /// Password-derived h scalar (w1 mod order).
    h: Scalar,
    /// Ephemeral private key x.
    x: Scalar,
    /// Our share L = x*G + w0*M
    share_p: ProjectivePoint,
    /// User random bytes (for context).
    user_random: [u8; 32],
    /// Device random bytes (from register response).
    dev_random: [u8; 32],
    /// Cipher suite.
    suite: Spake2PlusCipherSuite,
}

impl Spake2PlusProver {
    /// Create a new SPAKE2+ prover with derived password parameters.
    ///
    /// # Arguments
    ///
    /// * `password` - The password/credential bytes
    /// * `salt` - Salt from device (base64 decoded)
    /// * `iterations` - PBKDF2 iteration count
    /// * `user_random` - 32 random bytes generated by client
    /// * `dev_random` - 32 random bytes from device
    /// * `suite` - Cipher suite to use
    pub fn new(
        password: &[u8],
        salt: &[u8],
        iterations: u32,
        user_random: [u8; 32],
        dev_random: [u8; 32],
        suite: Spake2PlusCipherSuite,
    ) -> Self {
        // Derive w0 and w1 using PBKDF2
        let (w0, h) = derive_w0_w1(password, salt, iterations);

        // Generate ephemeral key x
        let x = generate_scalar();

        // Compute share: L = x*G + w0*M
        let m = decode_point(&P256_M_COMPRESSED);
        let share_p = ProjectivePoint::mul_by_generator(&x) + (m * w0);

        Self {
            w0,
            h,
            x,
            share_p,
            user_random,
            dev_random,
            suite,
        }
    }

    /// Returns our share (X = x*G + w0*M) as uncompressed SEC1 bytes.
    pub fn share(&self) -> Vec<u8> {
        let affine = self.share_p.to_affine();
        affine.to_encoded_point(false).as_bytes().to_vec()
    }

    /// Process the device's share and compute confirmation values.
    ///
    /// # Arguments
    ///
    /// * `dev_share_bytes` - Device's share Y (uncompressed SEC1 point)
    ///
    /// # Returns
    ///
    /// Tuple of (user_confirm, expected_dev_confirm, shared_key)
    #[allow(clippy::type_complexity)]
    pub fn process_share(
        &self,
        dev_share_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>, Vec<u8>), Spake2PlusError> {
        // Decode device share Y
        let y_point = decode_point_from_bytes(dev_share_bytes)?;

        // Decode N point
        let n = decode_point(&P256_N_COMPRESSED);

        // Compute R' = Y - w0*N
        let r_prime = y_point - (n * self.w0);

        // Compute Z = x * R' and V = h * R'
        let z = r_prime * self.x;
        let v = r_prime * self.h;

        // Encode points for transcript
        let m = decode_point(&P256_M_COMPRESSED);
        let m_enc = point_to_uncompressed(&m);
        let n_enc = point_to_uncompressed(&n);
        let l_enc = self.share(); // Our share
        let r_enc = dev_share_bytes.to_vec(); // Device share
        let z_enc = point_to_uncompressed(&z);
        let v_enc = point_to_uncompressed(&v);

        // Compute context hash
        let context_hash = self.compute_context_hash();

        // Encode w0 for transcript
        let w0_bytes = scalar_to_minimal_bytes(&self.w0);

        // Build transcript TT
        let transcript = build_transcript(
            &context_hash,
            &[], // idProver - empty
            &[], // idVerifier - empty
            &m_enc,
            &n_enc,
            &l_enc,
            &r_enc,
            &z_enc,
            &v_enc,
            &w0_bytes,
        );

        // Hash transcript to get K_main
        let k_main = self.hash(&transcript);

        // Derive confirmation keys
        let mac_len = self.suite.mac_len();
        let conf_keys = self.hkdf_expand("ConfirmationKeys", &k_main, mac_len * 2);
        let kc_a = &conf_keys[..mac_len];
        let kc_b = &conf_keys[mac_len..mac_len * 2];

        // Compute confirmations
        let user_confirm = self.compute_mac(kc_a, &r_enc);
        let expected_dev_confirm = self.compute_mac(kc_b, &l_enc);

        // Derive shared key
        let shared_key = self.hkdf_expand("SharedKey", &k_main, self.suite.digest_len());

        Ok((user_confirm, expected_dev_confirm, shared_key))
    }

    /// Compute the context hash: hash(PAKE_CONTEXT_TAG || user_random || dev_random)
    fn compute_context_hash(&self) -> Vec<u8> {
        let mut data = Vec::with_capacity(PAKE_CONTEXT_TAG.len() + 64);
        data.extend_from_slice(PAKE_CONTEXT_TAG);
        data.extend_from_slice(&self.user_random);
        data.extend_from_slice(&self.dev_random);
        self.hash(&data)
    }

    /// Hash data using the suite's hash algorithm.
    fn hash(&self, data: &[u8]) -> Vec<u8> {
        match self.suite {
            Spake2PlusCipherSuite::P256Sha256HmacSha256
            | Spake2PlusCipherSuite::P256Sha256CmacAes => {
                let mut hasher = Sha256::new();
                hasher.update(data);
                hasher.finalize().to_vec()
            }
            Spake2PlusCipherSuite::P256Sha512HmacSha512
            | Spake2PlusCipherSuite::P256Sha512CmacAes => {
                let mut hasher = Sha512::new();
                hasher.update(data);
                hasher.finalize().to_vec()
            }
        }
    }

    /// HKDF-Expand to derive key material.
    fn hkdf_expand(&self, info: &str, prk: &[u8], length: usize) -> Vec<u8> {
        match self.suite {
            Spake2PlusCipherSuite::P256Sha256HmacSha256
            | Spake2PlusCipherSuite::P256Sha256CmacAes => {
                hkdf_expand_sha256(prk, info.as_bytes(), length)
            }
            Spake2PlusCipherSuite::P256Sha512HmacSha512
            | Spake2PlusCipherSuite::P256Sha512CmacAes => {
                hkdf_expand_sha512(prk, info.as_bytes(), length)
            }
        }
    }

    /// Compute MAC for confirmation.
    fn compute_mac(&self, key: &[u8], data: &[u8]) -> Vec<u8> {
        if self.suite.uses_cmac() {
            cmac_aes(key, data)
        } else {
            match self.suite {
                Spake2PlusCipherSuite::P256Sha256HmacSha256 => hmac_sha256(key, data),
                Spake2PlusCipherSuite::P256Sha512HmacSha512 => hmac_sha512(key, data),
                _ => hmac_sha256(key, data),
            }
        }
    }
}

/// Error type for SPAKE2+ operations.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Spake2PlusError {
    /// Invalid point encoding.
    InvalidPoint,
    /// Confirmation mismatch.
    ConfirmationMismatch,
    /// Invalid parameters.
    InvalidParameters(String),
}

impl std::fmt::Display for Spake2PlusError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidPoint => write!(f, "invalid elliptic curve point"),
            Self::ConfirmationMismatch => write!(f, "SPAKE2+ confirmation mismatch"),
            Self::InvalidParameters(s) => write!(f, "invalid parameters: {}", s),
        }
    }
}

impl std::error::Error for Spake2PlusError {}

// =============================================================================
// Helper functions
// =============================================================================

/// Derive w0 and w1 from password using PBKDF2-SHA256.
fn derive_w0_w1(password: &[u8], salt: &[u8], iterations: u32) -> (Scalar, Scalar) {
    use pbkdf2::pbkdf2_hmac;

    // We need (256 + 64) * 2 bits = 640 bits = 80 bytes total
    // to have at least 64 bits of bias reduction
    let mut output = [0u8; 80];
    pbkdf2_hmac::<Sha256>(password, salt, iterations, &mut output);

    // First 40 bytes for w0, next 40 bytes for w1
    let w0_bytes = &output[..40];
    let w1_bytes = &output[40..80];

    // Reduce mod order
    let w0 = bytes_to_scalar_wide(w0_bytes);
    let w1 = bytes_to_scalar_wide(w1_bytes);

    (w0, w1)
}

/// Convert wide bytes to a scalar (reduced mod order).
fn bytes_to_scalar_wide(bytes: &[u8]) -> Scalar {
    // Convert to big integer and reduce mod order
    // P-256 order is: 0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551
    use p256::U256;
    use p256::elliptic_curve::ops::Reduce;

    // Hash down to 32 bytes if input is longer
    let hash = if bytes.len() > 32 {
        let mut hasher = Sha256::new();
        hasher.update(bytes);
        hasher.finalize().to_vec()
    } else {
        let mut padded = vec![0u8; 32];
        let start = 32 - bytes.len();
        padded[start..].copy_from_slice(bytes);
        padded
    };

    let wide = U256::from_be_slice(&hash);
    Scalar::reduce(wide)
}

/// Generate a random scalar.
fn generate_scalar() -> Scalar {
    use p256::U256;
    use p256::elliptic_curve::ops::Reduce;
    use rand::RngCore;

    let mut bytes = [0u8; 32];
    rand::rng().fill_bytes(&mut bytes);
    let n = U256::from_be_slice(&bytes);
    Scalar::reduce(n)
}

/// Decode a compressed SEC1 point to ProjectivePoint.
fn decode_point(compressed: &[u8; 33]) -> ProjectivePoint {
    let encoded = p256::EncodedPoint::from_bytes(compressed).expect("valid compressed point");
    let affine = AffinePoint::from_encoded_point(&encoded).expect("valid point on curve");
    ProjectivePoint::from(affine)
}

/// Decode SEC1 point bytes (compressed or uncompressed).
fn decode_point_from_bytes(bytes: &[u8]) -> Result<ProjectivePoint, Spake2PlusError> {
    let encoded =
        p256::EncodedPoint::from_bytes(bytes).map_err(|_| Spake2PlusError::InvalidPoint)?;
    let affine_opt = AffinePoint::from_encoded_point(&encoded);
    if affine_opt.is_none().into() {
        return Err(Spake2PlusError::InvalidPoint);
    }
    Ok(ProjectivePoint::from(affine_opt.unwrap()))
}

/// Convert ProjectivePoint to uncompressed SEC1 bytes.
fn point_to_uncompressed(point: &ProjectivePoint) -> Vec<u8> {
    let affine = point.to_affine();
    affine.to_encoded_point(false).as_bytes().to_vec()
}

/// Convert scalar to minimal big-endian bytes (as per RFC 9383).
fn scalar_to_minimal_bytes(scalar: &Scalar) -> Vec<u8> {
    let bytes = scalar.to_bytes();
    // Find first non-zero byte
    let first_nonzero = bytes.iter().position(|&b| b != 0).unwrap_or(31);
    let minimal = &bytes[first_nonzero..];

    // If leading bit is set, prepend 0x00
    if !minimal.is_empty() && (minimal[0] & 0x80) != 0 {
        let mut result = vec![0x00];
        result.extend_from_slice(minimal);
        result
    } else if minimal.is_empty() {
        vec![0x00]
    } else {
        minimal.to_vec()
    }
}

/// Build the SPAKE2+ transcript.
#[allow(clippy::too_many_arguments)]
fn build_transcript(
    context: &[u8],
    id_prover: &[u8],
    id_verifier: &[u8],
    m: &[u8],
    n: &[u8],
    share_p: &[u8],
    share_v: &[u8],
    z: &[u8],
    v: &[u8],
    w0: &[u8],
) -> Vec<u8> {
    let mut tt = Vec::new();

    // len8le encoding: 8-byte little-endian length prefix
    fn len8le(data: &[u8]) -> Vec<u8> {
        let mut result = Vec::with_capacity(8 + data.len());
        result.extend_from_slice(&(data.len() as u64).to_le_bytes());
        result.extend_from_slice(data);
        result
    }

    tt.extend(len8le(context));
    tt.extend(len8le(id_prover));
    tt.extend(len8le(id_verifier));
    tt.extend(len8le(m));
    tt.extend(len8le(n));
    tt.extend(len8le(share_p));
    tt.extend(len8le(share_v));
    tt.extend(len8le(z));
    tt.extend(len8le(v));
    tt.extend(len8le(w0));

    tt
}

/// HKDF-Expand using SHA-256.
fn hkdf_expand_sha256(prk: &[u8], info: &[u8], length: usize) -> Vec<u8> {
    use hkdf::Hkdf;

    // Use zero salt for expansion
    let salt = vec![0u8; 32];
    let hkdf = Hkdf::<Sha256>::new(Some(&salt), prk);
    let mut output = vec![0u8; length];
    hkdf.expand(info, &mut output)
        .expect("valid HKDF expansion");
    output
}

/// HKDF-Expand using SHA-512.
fn hkdf_expand_sha512(prk: &[u8], info: &[u8], length: usize) -> Vec<u8> {
    use hkdf::Hkdf;

    let salt = vec![0u8; 64];
    let hkdf = Hkdf::<Sha512>::new(Some(&salt), prk);
    let mut output = vec![0u8; length];
    hkdf.expand(info, &mut output)
        .expect("valid HKDF expansion");
    output
}

/// HMAC-SHA256.
fn hmac_sha256(key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut mac = Hmac::<Sha256>::new_from_slice(key).expect("valid key length");
    mac.update(data);
    mac.finalize().into_bytes().to_vec()
}

/// HMAC-SHA512.
fn hmac_sha512(key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut mac = Hmac::<Sha512>::new_from_slice(key).expect("valid key length");
    mac.update(data);
    mac.finalize().into_bytes().to_vec()
}

/// CMAC-AES-128.
fn cmac_aes(key: &[u8], data: &[u8]) -> Vec<u8> {
    use aes::Aes128;
    use cmac::{Cmac, Mac};

    // CMAC requires 16-byte key for AES-128
    let key_bytes: [u8; 16] = key[..16].try_into().expect("key at least 16 bytes");
    let mut mac = Cmac::<Aes128>::new_from_slice(&key_bytes).expect("valid key");
    mac.update(data);
    mac.finalize().into_bytes().to_vec()
}

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

    #[test]
    fn test_decode_m_point() {
        let m = decode_point(&P256_M_COMPRESSED);
        let encoded = point_to_uncompressed(&m);
        assert_eq!(encoded.len(), 65);
        assert_eq!(encoded[0], 0x04); // Uncompressed prefix
    }

    #[test]
    fn test_decode_n_point() {
        let n = decode_point(&P256_N_COMPRESSED);
        let encoded = point_to_uncompressed(&n);
        assert_eq!(encoded.len(), 65);
        assert_eq!(encoded[0], 0x04);
    }

    #[test]
    fn test_derive_w0_w1() {
        use p256::elliptic_curve::Field;

        let password = b"test_password";
        let salt = b"test_salt";
        let iterations = 1000;

        let (w0, w1) = derive_w0_w1(password, salt, iterations);

        // Verify scalars are non-zero
        assert!(!bool::from(w0.is_zero()));
        assert!(!bool::from(w1.is_zero()));
    }

    #[test]
    fn test_scalar_to_minimal_bytes() {
        let scalar = Scalar::ONE;
        let bytes = scalar_to_minimal_bytes(&scalar);
        assert_eq!(bytes, vec![0x01]);
    }

    #[test]
    fn test_hkdf_expand_sha256() {
        let prk = [0x42u8; 32];
        let info = b"test";
        let output = hkdf_expand_sha256(&prk, info, 16);
        assert_eq!(output.len(), 16);
    }
}
