use coracle_lib::encryption::{nip44, shared_secret, EncryptionError};
use coracle_lib::keys::{PublicKey, SecretKey};

fn sk(hex: &str) -> SecretKey {
    SecretKey::from_hex(hex).unwrap()
}

fn pk(sk: &SecretKey) -> PublicKey {
    sk.public_key()
}

#[test]
fn shared_secret_is_symmetric() {
    let a = SecretKey::generate();
    let b = SecretKey::generate();
    let pa = pk(&a);
    let pb = pk(&b);
    assert_eq!(shared_secret(&a, &pb), shared_secret(&b, &pa));
}

#[test]
fn nip44_roundtrip_short() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    let msg = "hi";
    let ct = nip44::encrypt(&alice, &pk(&bob), msg).unwrap();
    let pt = nip44::decrypt(&bob, &pk(&alice), &ct).unwrap();
    assert_eq!(pt, msg);
}

#[test]
fn nip44_roundtrip_long() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    let msg: String = "lorem ipsum ".repeat(500);
    let ct = nip44::encrypt(&alice, &pk(&bob), &msg).unwrap();
    let pt = nip44::decrypt(&bob, &pk(&alice), &ct).unwrap();
    assert_eq!(pt, msg);
}

#[test]
fn nip44_methods_on_secret_key() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    let ct = alice.nip44_encrypt(&pk(&bob), "via method").unwrap();
    let pt = bob.nip44_decrypt(&pk(&alice), &ct).unwrap();
    assert_eq!(pt, "via method");
}

#[test]
fn nip44_conversation_key_reuse() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    let ck_a = nip44::ConversationKey::derive(&alice, &pk(&bob));
    let ck_b = nip44::ConversationKey::derive(&bob, &pk(&alice));
    assert_eq!(ck_a.as_bytes(), ck_b.as_bytes());

    let ct = ck_a.encrypt("reusable").unwrap();
    let pt = ck_b.decrypt(&ct).unwrap();
    assert_eq!(pt, "reusable");
}

#[test]
fn nip44_rejects_empty_message() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    assert_eq!(
        nip44::encrypt(&alice, &pk(&bob), ""),
        Err(EncryptionError::MessageEmpty)
    );
}

#[test]
fn nip44_rejects_tampered_mac() {
    let alice = SecretKey::generate();
    let bob = SecretKey::generate();
    let ct = nip44::encrypt(&alice, &pk(&bob), "tamper me").unwrap();
    // Flip a bit near the end (inside the MAC region).
    let mut bytes = base64::Engine::decode(
        &base64::engine::general_purpose::STANDARD,
        ct.as_bytes(),
    )
    .unwrap();
    let last = bytes.len() - 1;
    bytes[last] ^= 0x01;
    let tampered =
        base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bytes);
    assert_eq!(
        nip44::decrypt(&bob, &pk(&alice), &tampered),
        Err(EncryptionError::InvalidMac)
    );
}

/// Official NIP-44 v2 test vector: sec1=1, sec2=2, nonce=0..01, plaintext="a".
#[test]
fn nip44_official_vector() {
    let sec1 = sk("0000000000000000000000000000000000000000000000000000000000000001");
    let sec2 = sk("0000000000000000000000000000000000000000000000000000000000000002");
    let expected_ck =
        hex::decode("c41c775356fd92eadc63ff5a0dc1da211b268cbea22316767095b2871ea1412d").unwrap();
    let expected_ct = "AgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABee0G5VSK0/9YypIObAtDKfYEAjD35uVkHyB0F4DwrcNaCXlCWZKaArsGrY6M9wnuTMxWfp1RTN9Xga8no+kF5Vsb";

    let ck = nip44::ConversationKey::derive(&sec1, &pk(&sec2));
    assert_eq!(ck.as_bytes(), expected_ck.as_slice());

    let mut nonce = [0u8; 32];
    nonce[31] = 1;
    let ct = ck.encrypt_with_nonce("a", &nonce).unwrap();
    assert_eq!(ct, expected_ct);

    let pt = ck.decrypt(&ct).unwrap();
    assert_eq!(pt, "a");
}