rustuya 0.3.0-rc.3

use crate::crypto::TuyaCipher;
use crate::error::{Result, TuyaError};
use crate::protocol::{
    CommandType, TuyaProtocol, Version, apply_update_dps, create_base_payload,
    legacy_lan_ext_stream, strip_status_heartbeat,
};
use base64::{Engine as _, engine::general_purpose};
use log::trace;
use md5::{Digest, Md5};
use serde_json::Value;

pub struct ProtocolV31;

impl TuyaProtocol for ProtocolV31 {
    fn version(&self) -> Version {
        Version::V3_1
    }

    fn get_effective_command(&self, command: CommandType) -> u32 {
        command as u32
    }

    fn generate_payload(
        &self,
        device_id: &str,
        command: CommandType,
        data: Option<Value>,
        cid: Option<&str>,
        t: u64,
    ) -> Result<(u32, Value)> {
        let cmd_to_send = self.get_effective_command(command);
        let mut payload =
            create_base_payload(device_id, cid, data.clone(), Some(t.to_string().into()));

        match command {
            CommandType::UpdateDps => apply_update_dps(&mut payload, data),
            CommandType::Control | CommandType::ControlNew | CommandType::DpQueryNew => {
                payload.remove("gwId");
            }
            CommandType::DpQuery => {
                // v3.1 keeps gwId/devId/uid/cid/t/dps unchanged.
            }
            CommandType::LanExtStream => {
                payload = legacy_lan_ext_stream(data);
            }
            CommandType::Status | CommandType::HeartBeat => strip_status_heartbeat(&mut payload),
            _ => {
                // Default: gwId, devId, uid, cid, t, dps
            }
        }

        let payload_obj = Value::Object(payload);
        trace!("v3.1 generated payload (cmd {cmd_to_send}): {payload_obj}");

        Ok((cmd_to_send, payload_obj))
    }

    fn pack_payload(&self, payload: &[u8], cmd: u32, cipher: &TuyaCipher) -> Result<Vec<u8>> {
        if cmd == CommandType::Control as u32 || cmd == CommandType::ControlNew as u32 {
            // 1. AES-128-ECB encrypt
            let encrypted = cipher.encrypt(payload, false, None, None, true)?;

            // 2. Base64 encode
            let b64_payload = general_purpose::STANDARD.encode(&encrypted);
            let b64_bytes = b64_payload.as_bytes();

            // 3. Compute MD5 over: b"data=" + base64payload + b"||lpv=3.1||" + local_key
            let mut hasher = Md5::new();
            hasher.update(b"data=");
            hasher.update(b64_bytes);
            hasher.update(b"||lpv=3.1||");
            hasher.update(cipher.key());

            let hash = hasher.finalize();
            let md5_hex = crate::crypto::hex_encode(&hash);
            let md5_slice = &md5_hex[8..24];

            // 4. Final payload: b"3.1" + md5slice + base64payload
            let mut final_payload = Vec::with_capacity(3 + 16 + b64_bytes.len());
            final_payload.extend_from_slice(b"3.1");
            final_payload.extend_from_slice(md5_slice.as_bytes());
            final_payload.extend_from_slice(b64_bytes);

            Ok(final_payload)
        } else {
            Ok(payload.to_vec())
        }
    }

    fn decrypt_payload(&self, payload: Vec<u8>, cipher: &TuyaCipher) -> Result<Vec<u8>> {
        if payload.starts_with(b"3.1") && payload.len() > 19 {
            // Strip "3.1" (3 bytes) and MD5 slice (16 bytes)
            let encrypted_b64 = &payload[19..];

            // Base64 decode
            let encrypted = general_purpose::STANDARD
                .decode(encrypted_b64)
                .map_err(|_| TuyaError::DecryptionFailed)?;

            // AES-ECB decrypt
            cipher.decrypt(&encrypted, false, None, None, None)
        } else {
            Ok(payload)
        }
    }

    fn has_version_header(&self, _payload: &[u8]) -> bool {
        false
    }

    fn requires_session_key(&self) -> bool {
        false
    }

    fn encrypt_session_key(
        &self,
        session_key: &[u8],
        cipher: &TuyaCipher,
        _nonce: &[u8],
    ) -> Result<Vec<u8>> {
        cipher.encrypt(session_key, false, None, None, false)
    }

    fn get_prefix(&self) -> u32 {
        crate::protocol::PREFIX_55AA
    }

    fn get_hmac_key<'a>(&self, _cipher_key: &'a [u8]) -> Option<&'a [u8]> {
        None
    }

    fn is_empty_payload_allowed(&self, _cmd: u32) -> bool {
        false
    }

    fn should_check_dev22_fallback(&self) -> bool {
        false
    }
}