use aes::{Aes128, Aes256};
use anyhow::{Context, Result};
use cbc::Decryptor;
use cipher::{block_padding::NoPadding, BlockDecryptMut, KeyIvInit, StreamCipher};

const MASTER_KEY: [u8; 32] = *include_bytes!("../assets/master.key.bin");

/// Split a TIDAL `keyId` security token into the (key, nonce) pair.
pub fn decrypt_security_token(token: &str) -> Result<([u8; 16], [u8; 8])> {
    use base64::{engine::general_purpose::STANDARD, Engine as _};
    let bytes = STANDARD.decode(token).context("security token base64")?;
    anyhow::ensure!(bytes.len() > 16 && (bytes.len() - 16) % 16 == 0, "bad security token length");
    let (iv, cipher) = bytes.split_at(16);
    let mut buf = cipher.to_vec();
    Decryptor::<Aes256>::new_from_slices(&MASTER_KEY, iv)
        .map_err(|e| anyhow::anyhow!("key/iv: {e}"))?
        .decrypt_padded_mut::<NoPadding>(&mut buf)
        .map_err(|e| anyhow::anyhow!("security token decrypt: {e}"))?;
    anyhow::ensure!(buf.len() >= 24, "decrypted token too short");
    Ok((buf[..16].try_into().unwrap(), buf[16..24].try_into().unwrap()))
}

/// AES-128-CTR decrypt (nonce || big-endian counter, matching pycryptodome's
/// `Counter.new(64, prefix=nonce)`).
pub fn decrypt_bytes(data: Vec<u8>, key: &[u8; 16], nonce: &[u8; 8]) -> Vec<u8> {
    let mut iv = [0u8; 16];
    iv[..8].copy_from_slice(nonce);
    let mut cipher = ctr::Ctr128BE::<Aes128>::new_from_slices(key, &iv).expect("ctr key/iv");
    let mut buf = data;
    cipher.apply_keystream(&mut buf);
    buf
}

#[cfg(test)]
mod tests {
    use super::*;
    use cipher::{block_padding::NoPadding, BlockEncryptMut};

    #[test]
    fn security_token_round_trip() {
        let key = *b"0123456789abcdef";
        let nonce = *b"12345678";
        let mut plain = Vec::new();
        plain.extend_from_slice(&key);
        plain.extend_from_slice(&nonce);
        plain.extend_from_slice(&[0u8; 8]);
        let iv = [7u8; 16];
        let mut buf = plain.clone();
        cbc::Encryptor::<Aes256>::new_from_slices(&MASTER_KEY, &iv)
            .unwrap()
            .encrypt_padded_mut::<NoPadding>(&mut buf, plain.len())
            .unwrap();
        let mut token_raw = iv.to_vec();
        token_raw.extend_from_slice(&buf);
        use base64::{engine::general_purpose::STANDARD, Engine as _};
        let (k, n) = decrypt_security_token(&STANDARD.encode(token_raw)).unwrap();
        assert_eq!(k, key);
        assert_eq!(n, nonce);
    }

    #[test]
    fn ctr_round_trip() {
        let key = *b"0123456789abcdef";
        let nonce = *b"12345678";
        let data = b"flac audio bytes...".to_vec();
        let enc = {
            let mut iv = [0u8; 16];
            iv[..8].copy_from_slice(&nonce);
            let mut c = ctr::Ctr128BE::<Aes128>::new((&key).into(), (&iv).into());
            let mut b = data.clone();
            c.apply_keystream(&mut b);
            b
        };
        assert_eq!(decrypt_bytes(enc, &key, &nonce), data);
    }
}