use aes::{Aes128, cipher::{BlockDecrypt, BlockEncrypt, KeyInit, generic_array::GenericArray}};
use aes_gcm::{Aes128Gcm, Nonce, aead::{AeadInPlace, KeyInit as AeadKeyInit}};
use anyhow::{anyhow, bail, Context, Result};
use base64::{
    alphabet,
    engine::{
        general_purpose::{GeneralPurpose, GeneralPurposeConfig, STANDARD},
        DecodePaddingMode,
    },
    Engine,
};

/// Shared discovery/bind key used by the original AES-128-ECB protocol.
pub const GENERIC_GREE_V1_KEY: &str = "a3K8Bx%2r8Y7#xDh";
/// Shared discovery/bind key used by AES-128-GCM capable Wi-Fi modules.
pub const GENERIC_GREE_V2_KEY: &str = "{yxAHAY_Lm6pbC/<";
/// GREE protocol v2 uses a fixed nonce and AAD, matching the EWPE/GREE LAN protocol.
const GCM_NONCE: [u8; 12] = [0x54, 0x40, 0x78, 0x44, 0x49, 0x67, 0x5a, 0x51, 0x6c, 0x5e, 0x63, 0x13];
const GCM_AAD: &[u8] = b"qualcomm-test";

// Some GREE Wi-Fi modules emit technically non-canonical Base64: padding may
// be omitted and unused trailing bits can be set. Python's base64 decoder,
// used by the established GREE implementations, accepts those packets. Keep
// encoding canonical but use a forgiving decoder for device-originated data.
const GREE_BASE64_DECODE: GeneralPurpose = GeneralPurpose::new(
    &alphabet::STANDARD,
    GeneralPurposeConfig::new()
        .with_decode_padding_mode(DecodePaddingMode::Indifferent)
        .with_decode_allow_trailing_bits(true),
);

pub fn normalize_key(key: &str) -> Result<[u8; 16]> {
    let bytes = key.as_bytes();
    if bytes.len() == 16 {
        let mut out = [0_u8; 16];
        out.copy_from_slice(bytes);
        return Ok(out);
    }
    if let Ok(decoded) = STANDARD.decode(key) {
        if decoded.len() == 16 {
            let mut out = [0_u8; 16];
            out.copy_from_slice(&decoded);
            return Ok(out);
        }
    }
    bail!("GREE key must contain 16 bytes or base64-encoded 16 bytes")
}

pub fn encrypt_v1(key: &str, plaintext: &[u8]) -> Result<String> {
    let key = normalize_key(key)?;
    let cipher = Aes128::new_from_slice(&key).map_err(|_| anyhow!("invalid AES key"))?;
    let pad = 16 - (plaintext.len() % 16);
    let mut data = Vec::with_capacity(plaintext.len() + pad);
    data.extend_from_slice(plaintext);
    data.extend(std::iter::repeat(pad as u8).take(pad));
    for block in data.chunks_exact_mut(16) {
        cipher.encrypt_block(GenericArray::from_mut_slice(block));
    }
    Ok(STANDARD.encode(data))
}

pub fn decrypt_v1(key: &str, ciphertext_b64: &str) -> Result<Vec<u8>> {
    let key = normalize_key(key)?;
    let cipher = Aes128::new_from_slice(&key).map_err(|_| anyhow!("invalid AES key"))?;
    let mut data = GREE_BASE64_DECODE.decode(ciphertext_b64).context("invalid base64 packet")?;
    if data.is_empty() || data.len() % 16 != 0 {
        bail!("invalid AES-ECB ciphertext length")
    }
    for block in data.chunks_exact_mut(16) {
        cipher.decrypt_block(GenericArray::from_mut_slice(block));
    }
    let pad = *data.last().ok_or_else(|| anyhow!("empty plaintext"))? as usize;
    let valid_padding = pad > 0
        && pad <= 16
        && data.len() >= pad
        && data[data.len() - pad..].iter().all(|v| *v as usize == pad);
    if valid_padding {
        data.truncate(data.len() - pad);
    } else if let Some(last_json_byte) = data.iter().rposition(|byte| *byte == b'}') {
        // Legacy GREE implementations are deliberately tolerant here: some
        // modules return non-standard padding but the JSON itself is valid.
        data.truncate(last_json_byte + 1);
    } else {
        bail!("invalid AES-ECB padding and no JSON terminator")
    }
    Ok(data)
}

#[derive(Debug, Clone)]
pub struct V2Encrypted {
    pub ciphertext: String,
    pub tag: String,
}

pub fn encrypt_v2(key: &str, plaintext: &[u8]) -> Result<V2Encrypted> {
    let key = normalize_key(key)?;
    let cipher = <Aes128Gcm as AeadKeyInit>::new_from_slice(&key).map_err(|_| anyhow!("invalid AES-GCM key"))?;
    let nonce = Nonce::from_slice(&GCM_NONCE);
    let mut buffer = plaintext.to_vec();
    let tag = cipher.encrypt_in_place_detached(nonce, GCM_AAD, &mut buffer)
        .map_err(|_| anyhow!("AES-GCM encryption failed"))?;
    Ok(V2Encrypted {
        ciphertext: STANDARD.encode(buffer),
        tag: STANDARD.encode(tag),
    })
}

pub fn decrypt_v2(key: &str, ciphertext_b64: &str, tag_b64: &str) -> Result<Vec<u8>> {
    let key = normalize_key(key)?;
    let cipher = <Aes128Gcm as AeadKeyInit>::new_from_slice(&key).map_err(|_| anyhow!("invalid AES-GCM key"))?;
    let tag_bytes = GREE_BASE64_DECODE.decode(tag_b64).context("invalid GCM tag")?;
    if tag_bytes.len() != 16 { bail!("invalid GCM tag length: {} bytes", tag_bytes.len()) }
    let mut data = GREE_BASE64_DECODE.decode(ciphertext_b64).context("invalid GCM ciphertext")?;
    let nonce = Nonce::from_slice(&GCM_NONCE);
    let tag = GenericArray::from_slice(&tag_bytes);
    cipher.decrypt_in_place_detached(nonce, GCM_AAD, &mut data, tag)
        .map_err(|_| anyhow!("AES-GCM authentication failed"))?;
    // A few modules append 0xff filler bytes to decrypted JSON.
    data.retain(|byte| *byte != 0xff);
    Ok(data)
}

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

    #[test]
    fn v1_round_trip() {
        let value = br#"{"t":"status","mac":"112233445566"}"#;
        let encrypted = encrypt_v1(GENERIC_GREE_V1_KEY, value).unwrap();
        assert_eq!(decrypt_v1(GENERIC_GREE_V1_KEY, &encrypted).unwrap(), value);
    }

    #[test]
    fn v2_round_trip() {
        let value = b"gree-gcm-test";
        let encrypted = encrypt_v2(GENERIC_GREE_V2_KEY, value).unwrap();
        assert_eq!(decrypt_v2(GENERIC_GREE_V2_KEY, &encrypted.ciphertext, &encrypted.tag).unwrap(), value);
    }

    #[test]
    fn gree_base64_accepts_noncanonical_trailing_bits() {
        // 16 zero bytes canonically end with `A==`. `B==` carries the same
        // useful two bits but has non-zero unused trailing bits. Python's
        // base64.b64decode accepts it and real GREE modules emit this form.
        let decoded = GREE_BASE64_DECODE.decode("AAAAAAAAAAAAAAAAAAAAAB==").unwrap();
        assert_eq!(decoded, vec![0_u8; 16]);
    }
}