//! KLAP encryption for the TP-Link KLAP protocol.
//!
//! KLAP (Kasa Local Authentication Protocol) is used by newer firmware versions
//! on TP-Link devices. It uses AES-128-CBC encryption with keys derived from
//! a two-phase handshake.
//!
//! # Protocol Overview
//!
//! 1. **Authentication Hash**: `md5(md5(username) + md5(password))`
//! 2. **Handshake 1**: Client sends 16 random bytes, receives 16 bytes + hash
//! 3. **Handshake 2**: Client sends confirmation hash, receives session cookie
//! 4. **Encryption**: AES-128-CBC with keys derived from seeds and auth hash

use aes::Aes128;
use cbc::{
    Decryptor, Encryptor,
    cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit, block_padding::Pkcs7},
};
use md5::{Digest as Md5Digest, Md5};
use sha1::Sha1;
use sha2::Sha256;

use crate::Credentials;

type Aes128CbcEnc = Encryptor<Aes128>;
type Aes128CbcDec = Decryptor<Aes128>;

/// Generates the authentication hash for KLAP protocol (v1).
///
/// The auth hash is computed as: `md5(md5(username) + md5(password))`
///
/// # Arguments
///
/// * `credentials` - The user credentials
///
/// # Returns
///
/// A 16-byte MD5 hash used for authentication.
///
/// # Example
///
/// ```
/// use kasa_core::{Credentials, crypto::generate_auth_hash};
///
/// let creds = Credentials::new("user@example.com", "password123");
/// let hash = generate_auth_hash(&creds);
/// assert_eq!(hash.len(), 16);
/// ```
pub fn generate_auth_hash(credentials: &Credentials) -> [u8; 16] {
    let username_hash = md5_hash(credentials.username.as_bytes());
    let password_hash = md5_hash(credentials.password.as_bytes());

    let mut combined = Vec::with_capacity(32);
    combined.extend_from_slice(&username_hash);
    combined.extend_from_slice(&password_hash);

    md5_hash(&combined)
}

/// Generates the authentication hash for KLAP protocol v2.
///
/// The auth hash is computed as: `sha256(sha1(username) + sha1(password))`
///
/// This is used by newer devices (KlapTransportV2).
pub fn generate_auth_hash_v2(credentials: &Credentials) -> [u8; 32] {
    let username_hash = sha1_hash(credentials.username.as_bytes());
    let password_hash = sha1_hash(credentials.password.as_bytes());

    let mut combined = Vec::with_capacity(40);
    combined.extend_from_slice(&username_hash);
    combined.extend_from_slice(&password_hash);

    sha256_hash(&combined)
}

/// Generates the owner hash for KLAP protocol.
///
/// The owner hash is: `md5(username)`
///
/// This is used to identify the device owner.
///
/// # Arguments
///
/// * `credentials` - The user credentials
///
/// # Returns
///
/// A 16-byte MD5 hash of the username.
pub fn generate_owner_hash(credentials: &Credentials) -> [u8; 16] {
    md5_hash(credentials.username.as_bytes())
}

/// Computes the expected server response for handshake1 (v1).
///
/// The server response should be: `sha256(local_seed + auth_hash)`
///
/// Note: remote_seed is NOT included in the hash - it's only returned BY the device.
///
/// # Arguments
///
/// * `local_seed` - 16-byte random seed generated by the client
/// * `auth_hash` - 16-byte authentication hash
///
/// # Returns
///
/// A 32-byte SHA-256 hash that should match the server's response.
pub fn handshake1_seed_auth_hash(local_seed: &[u8], auth_hash: &[u8]) -> [u8; 32] {
    let mut data = Vec::with_capacity(local_seed.len() + auth_hash.len());
    data.extend_from_slice(local_seed);
    data.extend_from_slice(auth_hash);
    sha256_hash(&data)
}

/// Computes the expected server response for handshake1 (v2).
///
/// The server response should be: `sha256(local_seed + remote_seed + auth_hash)`
pub fn handshake1_seed_auth_hash_v2(
    local_seed: &[u8],
    remote_seed: &[u8],
    auth_hash: &[u8],
) -> [u8; 32] {
    let mut data = Vec::with_capacity(local_seed.len() + remote_seed.len() + auth_hash.len());
    data.extend_from_slice(local_seed);
    data.extend_from_slice(remote_seed);
    data.extend_from_slice(auth_hash);
    sha256_hash(&data)
}

/// Computes the client payload for handshake2 (v1).
///
/// The client sends: `sha256(remote_seed + auth_hash)`
///
/// Note: local_seed is NOT included.
///
/// # Arguments
///
/// * `remote_seed` - 16-byte random seed received from the device
/// * `auth_hash` - 16-byte authentication hash
///
/// # Returns
///
/// A 32-byte SHA-256 hash to send to the device.
pub fn handshake2_seed_auth_hash(remote_seed: &[u8], auth_hash: &[u8]) -> [u8; 32] {
    let mut data = Vec::with_capacity(remote_seed.len() + auth_hash.len());
    data.extend_from_slice(remote_seed);
    data.extend_from_slice(auth_hash);
    sha256_hash(&data)
}

/// Computes the client payload for handshake2 (v2).
///
/// The client sends: `sha256(remote_seed + local_seed + auth_hash)`
pub fn handshake2_seed_auth_hash_v2(
    local_seed: &[u8],
    remote_seed: &[u8],
    auth_hash: &[u8],
) -> [u8; 32] {
    let mut data = Vec::with_capacity(remote_seed.len() + local_seed.len() + auth_hash.len());
    data.extend_from_slice(remote_seed);
    data.extend_from_slice(local_seed);
    data.extend_from_slice(auth_hash);
    sha256_hash(&data)
}

/// Encryption session for KLAP protocol.
///
/// After a successful handshake, this session holds the encryption keys
/// and sequence number for encrypted communication.
#[derive(Debug, Clone)]
pub struct KlapEncryptionSession {
    /// AES key derived from seeds and auth hash.
    key: [u8; 16],
    /// Base IV (initialization vector) derived from seeds and auth hash (first 12 bytes).
    iv: [u8; 12],
    /// Signature prefix for payload verification.
    sig: [u8; 28],
    /// Current sequence number (incremented for each request).
    seq: i32,
}

impl KlapEncryptionSession {
    /// Creates a new encryption session from handshake results.
    ///
    /// # Arguments
    ///
    /// * `local_seed` - 16-byte client seed from handshake
    /// * `remote_seed` - 16-byte device seed from handshake
    /// * `auth_hash` - 16-byte authentication hash
    ///
    /// # Returns
    ///
    /// A new encryption session ready for use.
    pub fn new(local_seed: &[u8], remote_seed: &[u8], auth_hash: &[u8]) -> Self {
        // Key = first 16 bytes of sha256("lsk" + local_seed + remote_seed + auth_hash)
        let key = Self::derive_key(local_seed, remote_seed, auth_hash);

        // IV derivation: sha256("iv" + local_seed + remote_seed + auth_hash)
        // First 12 bytes are the IV base, last 4 bytes (of the full 16) are the initial seq
        let (iv, seq) = Self::derive_iv_and_seq(local_seed, remote_seed, auth_hash);

        // Signature = first 28 bytes of sha256("ldk" + local_seed + remote_seed + auth_hash)
        let sig = Self::derive_sig(local_seed, remote_seed, auth_hash);

        Self { key, iv, sig, seq }
    }

    /// Returns the current sequence number.
    pub fn seq(&self) -> i32 {
        self.seq
    }

    /// Encrypts a payload for sending to the device.
    ///
    /// The payload is encrypted using AES-128-CBC with PKCS7 padding.
    /// A signature hash is prepended.
    ///
    /// # Arguments
    ///
    /// * `plaintext` - The JSON command to encrypt
    ///
    /// # Returns
    ///
    /// The encrypted payload ready to send via HTTP POST.
    pub fn encrypt(&mut self, plaintext: &[u8]) -> Vec<u8> {
        self.seq = self.seq.wrapping_add(1);
        let iv = self.build_iv(self.seq);

        // Encrypt the payload with AES-128-CBC first
        let cipher = Aes128CbcEnc::new(&self.key.into(), &iv.into());

        // Calculate padded size (multiple of 16)
        let plaintext_len = plaintext.len();
        let padded_len = ((plaintext_len / 16) + 1) * 16;
        let mut buffer = vec![0u8; padded_len];
        buffer[..plaintext_len].copy_from_slice(plaintext);

        let encrypted = cipher
            .encrypt_padded_mut::<Pkcs7>(&mut buffer, plaintext_len)
            .expect("encryption should not fail with valid padding");

        // Create signature: sha256(sig + seq_bytes + ciphertext)
        let mut sig_data = Vec::with_capacity(28 + 4 + encrypted.len());
        sig_data.extend_from_slice(&self.sig);
        sig_data.extend_from_slice(&self.seq.to_be_bytes());
        sig_data.extend_from_slice(encrypted);
        let signature = sha256_hash(&sig_data);

        // Prepend signature to encrypted payload
        let mut result = Vec::with_capacity(32 + encrypted.len());
        result.extend_from_slice(&signature);
        result.extend_from_slice(encrypted);

        result
    }

    /// Decrypts a response from the device.
    ///
    /// # Arguments
    ///
    /// * `ciphertext` - The encrypted response (signature + encrypted payload)
    /// * `seq` - The sequence number used for this request
    ///
    /// # Returns
    ///
    /// The decrypted JSON response, or an error if decryption fails.
    pub fn decrypt(&self, ciphertext: &[u8], seq: i32) -> Result<Vec<u8>, DecryptError> {
        if ciphertext.len() < 32 {
            return Err(DecryptError::PayloadTooShort);
        }

        let _signature = &ciphertext[..32];
        let encrypted = &ciphertext[32..];

        let iv = self.build_iv(seq);

        // Decrypt with AES-128-CBC
        let cipher = Aes128CbcDec::new(&self.key.into(), &iv.into());
        let mut buffer = encrypted.to_vec();

        let decrypted = cipher
            .decrypt_padded_mut::<Pkcs7>(&mut buffer)
            .map_err(|_| DecryptError::DecryptionFailed)?;

        Ok(decrypted.to_vec())
    }

    /// Builds the full IV from base IV and sequence number.
    fn build_iv(&self, seq: i32) -> [u8; 16] {
        let mut iv = [0u8; 16];
        iv[..12].copy_from_slice(&self.iv);
        iv[12..16].copy_from_slice(&seq.to_be_bytes());
        iv
    }

    fn derive_key(local: &[u8], remote: &[u8], auth: &[u8]) -> [u8; 16] {
        let mut data = Vec::with_capacity(3 + local.len() + remote.len() + auth.len());
        data.extend_from_slice(b"lsk");
        data.extend_from_slice(local);
        data.extend_from_slice(remote);
        data.extend_from_slice(auth);
        let hash = sha256_hash(&data);
        let mut key = [0u8; 16];
        key.copy_from_slice(&hash[..16]);
        key
    }

    fn derive_iv_and_seq(local: &[u8], remote: &[u8], auth: &[u8]) -> ([u8; 12], i32) {
        let mut data = Vec::with_capacity(2 + local.len() + remote.len() + auth.len());
        data.extend_from_slice(b"iv");
        data.extend_from_slice(local);
        data.extend_from_slice(remote);
        data.extend_from_slice(auth);
        let hash = sha256_hash(&data);

        // First 12 bytes are IV base
        let mut iv = [0u8; 12];
        iv.copy_from_slice(&hash[..12]);

        // Last 4 bytes of the full IV (bytes 12-15 of hash) are initial sequence
        let seq = i32::from_be_bytes([hash[28], hash[29], hash[30], hash[31]]);

        (iv, seq)
    }

    fn derive_sig(local: &[u8], remote: &[u8], auth: &[u8]) -> [u8; 28] {
        let mut data = Vec::with_capacity(3 + local.len() + remote.len() + auth.len());
        data.extend_from_slice(b"ldk");
        data.extend_from_slice(local);
        data.extend_from_slice(remote);
        data.extend_from_slice(auth);
        let hash = sha256_hash(&data);
        let mut sig = [0u8; 28];
        sig.copy_from_slice(&hash[..28]);
        sig
    }
}

/// Error type for KLAP decryption failures.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecryptError {
    /// The encrypted payload is too short (less than 32 bytes for signature).
    PayloadTooShort,
    /// AES decryption or padding validation failed.
    DecryptionFailed,
    /// Signature verification failed.
    SignatureInvalid,
}

impl std::fmt::Display for DecryptError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DecryptError::PayloadTooShort => write!(f, "payload too short"),
            DecryptError::DecryptionFailed => write!(f, "decryption failed"),
            DecryptError::SignatureInvalid => write!(f, "signature invalid"),
        }
    }
}

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

/// Helper function to compute MD5 hash.
fn md5_hash(data: &[u8]) -> [u8; 16] {
    let mut hasher = Md5::new();
    hasher.update(data);
    hasher.finalize().into()
}

/// Helper function to compute SHA1 hash.
fn sha1_hash(data: &[u8]) -> [u8; 20] {
    let mut hasher = Sha1::new();
    hasher.update(data);
    hasher.finalize().into()
}

/// Helper function to compute SHA256 hash.
fn sha256_hash(data: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().into()
}

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

    #[test]
    fn test_generate_auth_hash() {
        // Test with known values
        let creds = Credentials::new("test@example.com", "password123");
        let hash = generate_auth_hash(&creds);
        assert_eq!(hash.len(), 16);

        // Same credentials should produce same hash
        let hash2 = generate_auth_hash(&creds);
        assert_eq!(hash, hash2);

        // Different credentials should produce different hash
        let creds2 = Credentials::new("other@example.com", "password123");
        let hash3 = generate_auth_hash(&creds2);
        assert_ne!(hash, hash3);
    }

    #[test]
    fn test_generate_auth_hash_blank() {
        let blank = Credentials::blank();
        let hash = generate_auth_hash(&blank);
        assert_eq!(hash.len(), 16);
    }

    #[test]
    fn test_handshake1_v1_only_uses_local_seed() {
        let local = [1u8; 16];
        let auth = [3u8; 16];

        let h1 = handshake1_seed_auth_hash(&local, &auth);
        assert_eq!(h1.len(), 32);
    }

    #[test]
    fn test_handshake2_v1_only_uses_remote_seed() {
        let remote = [2u8; 16];
        let auth = [3u8; 16];

        let h2 = handshake2_seed_auth_hash(&remote, &auth);
        assert_eq!(h2.len(), 32);
    }

    #[test]
    fn test_encryption_session_creation() {
        let local = [1u8; 16];
        let remote = [2u8; 16];
        let auth = [3u8; 16];

        let session = KlapEncryptionSession::new(&local, &remote, &auth);
        // Just verify it doesn't panic and has a sequence number
        let _ = session.seq();
    }

    #[test]
    fn test_encrypt_decrypt_roundtrip() {
        let local = [1u8; 16];
        let remote = [2u8; 16];
        let auth = [3u8; 16];

        let mut session = KlapEncryptionSession::new(&local, &remote, &auth);
        let seq_before = session.seq();

        let plaintext = b"test message";
        let encrypted = session.encrypt(plaintext);

        // Sequence should have incremented
        let seq_after = session.seq();
        assert_eq!(seq_after, seq_before.wrapping_add(1));

        // Decrypt with the same session (but use the seq that was used for encryption)
        let decrypted = session.decrypt(&encrypted, seq_after).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_decrypt_payload_too_short() {
        let local = [1u8; 16];
        let remote = [2u8; 16];
        let auth = [3u8; 16];

        let session = KlapEncryptionSession::new(&local, &remote, &auth);
        let result = session.decrypt(&[0u8; 16], 0);
        assert_eq!(result, Err(DecryptError::PayloadTooShort));
    }
}
