ex_reticulum

use rustler::{Binary, Encoder, Env, OwnedBinary, ResourceArc, Term};
use rand_core::OsRng;
use ed25519_dalek::Signature;
use reticulum::identity::{Identity, PrivateIdentity, DecryptIdentity, DerivedKey, PUBLIC_KEY_LENGTH};
use x25519_dalek::StaticSecret;

use crate::atoms;
use crate::resources::{IdentityKind, IdentityResource};

fn address_hash_binary<'a>(env: Env<'a>, kind: &IdentityKind) -> Option<Binary<'a>> {
    let hash_bytes = match kind {
        IdentityKind::Private(p) => p.address_hash().as_slice(),
        IdentityKind::Public(p) => p.address_hash.as_slice(),
    };
    let mut bin = OwnedBinary::new(hash_bytes.len())?;
    bin.as_mut_slice().copy_from_slice(hash_bytes);
    Some(Binary::from_owned(bin, env))
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_generate(env: Env) -> Term {
    let private = PrivateIdentity::new_from_rand(OsRng);
    let resource = ResourceArc::new(IdentityResource(
        std::sync::Mutex::new(IdentityKind::Private(private)),
    ));
    let hash = {
        let lock = match resource.0.lock() {
            Ok(l) => l,
            Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
        };
        match address_hash_binary(env, &lock) {
            Some(b) => b,
            None => return (atoms::error(), atoms::out_of_memory()).encode(env),
        }
    };
    (atoms::ok(), resource, hash).encode(env)
}

#[rustler::nif]
fn identity_address_hash(env: Env, identity: ResourceArc<IdentityResource>) -> Term {
    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };
    let hash = match address_hash_binary(env, &lock) {
        Some(b) => b,
        None => return (atoms::error(), atoms::out_of_memory()).encode(env),
    };
    (atoms::ok(), hash).encode(env)
}

#[rustler::nif]
fn identity_to_hex(env: Env, identity: ResourceArc<IdentityResource>) -> Term {
    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };
    let hex = match &*lock {
        IdentityKind::Private(p) => p.to_hex_string(),
        IdentityKind::Public(p) => p.to_hex_string(),
    };
    (atoms::ok(), hex).encode(env)
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_from_hex(env: Env, hex: String) -> Term {
    match PrivateIdentity::new_from_hex_string(&hex) {
        Ok(private) => {
            let resource = ResourceArc::new(IdentityResource(
                std::sync::Mutex::new(IdentityKind::Private(private)),
            ));
            let hash = {
                let lock = match resource.0.lock() {
                    Ok(l) => l,
                    Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
                };
                match address_hash_binary(env, &lock) {
                    Some(b) => b,
                    None => return (atoms::error(), atoms::out_of_memory()).encode(env),
                }
            };
            (atoms::ok(), resource, hash, atoms::private()).encode(env)
        }
        Err(_) => match Identity::new_from_hex_string(&hex) {
            Ok(public) => {
                let resource = ResourceArc::new(IdentityResource(
                    std::sync::Mutex::new(IdentityKind::Public(public)),
                ));
                let hash = {
                    let lock = match resource.0.lock() {
                        Ok(l) => l,
                        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
                    };
                    match address_hash_binary(env, &lock) {
                        Some(b) => b,
                        None => return (atoms::error(), atoms::out_of_memory()).encode(env),
                    }
                };
                (atoms::ok(), resource, hash, atoms::public()).encode(env)
            }
            Err(_) => (atoms::error(), atoms::invalid_argument()).encode(env),
        },
    }
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_from_name(env: Env, name: String) -> Term {
    let private = PrivateIdentity::new_from_name(&name);
    let resource = ResourceArc::new(IdentityResource(
        std::sync::Mutex::new(IdentityKind::Private(private)),
    ));
    let hash = {
        let lock = match resource.0.lock() {
            Ok(l) => l,
            Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
        };
        match address_hash_binary(env, &lock) {
            Some(b) => b,
            None => return (atoms::error(), atoms::out_of_memory()).encode(env),
        }
    };
    (atoms::ok(), resource, hash).encode(env)
}

#[rustler::nif]
fn identity_public_key(env: Env, identity: ResourceArc<IdentityResource>) -> Term {
    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };
    let key_bytes = match &*lock {
        IdentityKind::Private(p) => p.as_identity().public_key_bytes(),
        IdentityKind::Public(p) => p.public_key_bytes(),
    };
    let mut bin = match OwnedBinary::new(32) {
        Some(b) => b,
        None => return (atoms::error(), atoms::out_of_memory()).encode(env),
    };
    bin.as_mut_slice().copy_from_slice(key_bytes);
    (atoms::ok(), Binary::from_owned(bin, env)).encode(env)
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_sign<'a>(env: Env<'a>, identity: ResourceArc<IdentityResource>, data: Binary<'a>) -> Term<'a> {
    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };
    match &*lock {
        IdentityKind::Private(private) => {
            let sig = private.sign(data.as_slice());
            let sig_bytes = sig.to_bytes();
            let mut bin = match OwnedBinary::new(64) {
                Some(b) => b,
                None => return (atoms::error(), atoms::out_of_memory()).encode(env),
            };
            bin.as_mut_slice().copy_from_slice(&sig_bytes);
            (atoms::ok(), Binary::from_owned(bin, env)).encode(env)
        }
        IdentityKind::Public(_) => (atoms::error(), atoms::invalid_argument()).encode(env),
    }
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_verify<'a>(env: Env<'a>, identity: ResourceArc<IdentityResource>, data: Binary<'a>, sig: Binary<'a>) -> Term<'a> {
    if sig.as_slice().len() != 64 {
        return (atoms::error(), atoms::invalid_argument()).encode(env);
    }
    let signature = match Signature::from_slice(sig.as_slice()) {
        Ok(s) => s,
        Err(_) => return (atoms::error(), atoms::invalid_argument()).encode(env),
    };

    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };
    let result = match &*lock {
        IdentityKind::Private(private) => private.verify(data.as_slice(), &signature),
        IdentityKind::Public(public) => public.verify(data.as_slice(), &signature),
    };

    match result {
        Ok(()) => (atoms::ok(), atoms::verified()).encode(env),
        Err(_) => (atoms::error(), atoms::incorrect_signature()).encode(env),
    }
}

#[rustler::nif]
fn identity_from_public_keys<'a>(env: Env<'a>, pub_key: Binary<'a>, ver_key: Binary<'a>) -> Term<'a> {
    if pub_key.as_slice().len() != 32 || ver_key.as_slice().len() != 32 {
        return (atoms::error(), atoms::invalid_argument()).encode(env);
    }
    let identity = Identity::new_from_slices(pub_key.as_slice(), ver_key.as_slice());
    let resource = ResourceArc::new(IdentityResource(
        std::sync::Mutex::new(IdentityKind::Public(identity)),
    ));
    let hash = {
        let lock = match resource.0.lock() {
            Ok(l) => l,
            Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
        };
        match address_hash_binary(env, &lock) {
            Some(b) => b,
            None => return (atoms::error(), atoms::out_of_memory()).encode(env),
        }
    };
    (atoms::ok(), resource, hash).encode(env)
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_encrypt<'a>(
    env: Env<'a>,
    recipient: ResourceArc<IdentityResource>,
    plaintext: Binary<'a>,
) -> Term<'a> {
    let lock = match recipient.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };

    // Extract the recipient's public key regardless of identity type
    let recipient_pub_key = match &*lock {
        IdentityKind::Public(identity) => identity.public_key,
        IdentityKind::Private(private) => private.as_identity().public_key,
    };
    drop(lock);

    let text = plaintext.as_slice();
    // Output: 32 (ephemeral pubkey) + plaintext + 48 (fernet: 16 IV + 32 HMAC) + 16 (padding)
    let mut out_buf = vec![0u8; text.len() + 128];

    // Create ephemeral X25519 keypair, derive shared key via ECDH, Fernet-encrypt
    let ephemeral_secret = StaticSecret::random_from_rng(OsRng);
    let ephemeral_public = x25519_dalek::PublicKey::from(&ephemeral_secret);

    // Prepend ephemeral public key so recipient can reconstruct shared secret
    out_buf[..PUBLIC_KEY_LENGTH].copy_from_slice(ephemeral_public.as_bytes());

    let derived_key = DerivedKey::new_from_private_key(
        &ephemeral_secret,
        &recipient_pub_key,
        None,
    );

    use reticulum::crypt::fernet::{Fernet, PlainText};
    let fernet = Fernet::new_from_slices(
        &derived_key.as_bytes()[..derived_key.as_bytes().len() / 2],
        &derived_key.as_bytes()[derived_key.as_bytes().len() / 2..],
        OsRng,
    );

    let result = match fernet.encrypt(PlainText::from(text), &mut out_buf[PUBLIC_KEY_LENGTH..]) {
        Ok(token) => {
            let total_len = PUBLIC_KEY_LENGTH + token.as_bytes().len();
            Ok(&out_buf[..total_len])
        }
        Err(e) => Err(e),
    };

    match result {
        Ok(encrypted) => {
            let mut bin = match OwnedBinary::new(encrypted.len()) {
                Some(b) => b,
                None => return (atoms::error(), atoms::out_of_memory()).encode(env),
            };
            bin.as_mut_slice().copy_from_slice(encrypted);
            (atoms::ok(), Binary::from_owned(bin, env)).encode(env)
        }
        Err(_) => (atoms::error(), atoms::encryption_error()).encode(env),
    }
}

#[rustler::nif(schedule = "DirtyCpu")]
fn identity_decrypt<'a>(
    env: Env<'a>,
    identity: ResourceArc<IdentityResource>,
    ciphertext: Binary<'a>,
) -> Term<'a> {
    let lock = match identity.0.lock() {
        Ok(l) => l,
        Err(_) => return (atoms::error(), atoms::lock_error()).encode(env),
    };

    let private = match &*lock {
        IdentityKind::Private(p) => p,
        IdentityKind::Public(_) => return (atoms::error(), atoms::invalid_argument()).encode(env),
    };

    let data = ciphertext.as_slice();
    if data.len() < PUBLIC_KEY_LENGTH {
        return (atoms::error(), atoms::invalid_argument()).encode(env);
    }

    // Extract ephemeral public key from first 32 bytes
    let ephemeral_key = x25519_dalek::PublicKey::from(
        <[u8; 32]>::try_from(&data[..PUBLIC_KEY_LENGTH]).unwrap(),
    );
    let derived_key = private.derive_key(&ephemeral_key, None);

    let mut out_buf = vec![0u8; data.len()];

    match private.decrypt(OsRng, &data[PUBLIC_KEY_LENGTH..], &derived_key, &mut out_buf) {
        Ok(decrypted) => {
            let mut bin = match OwnedBinary::new(decrypted.len()) {
                Some(b) => b,
                None => return (atoms::error(), atoms::out_of_memory()).encode(env),
            };
            bin.as_mut_slice().copy_from_slice(decrypted);
            (atoms::ok(), Binary::from_owned(bin, env)).encode(env)
        }
        Err(_) => (atoms::error(), atoms::encryption_error()).encode(env),
    }
}