module cryptography

import x.crypto.ascon
import crypto.rand
import crypto.hmac
import x.crypto.curve25519
import crypto.sha256
import encoding.base64
import structs

fn ascon_key_generator() string {
	key := rand.bytes(ascon.key_size) or {
        logs.error('cryptography:cryptography:ascon_key_generator', 'ascon key generation failed due to: \'$err\'')
        return ''
    }
	return key.hex()
}

fn ascon_encrypt(key_hex string, plaintext string) string {
    key := []u8(key_hex.hex())
    nonce := rand.bytes(ascon.nonce_size) or {
        logs.error('cryptography:cryptography:ascon_encrypt', 'nonce generation failed due to: \'$err\'')
        return ''
    }
    ad := []u8{}

    ciphertext := ascon.encrypt(key, nonce, ad, plaintext.bytes()) or {
        logs.error('cryptography:cryptography:ascon_encrypt', 'encrypt failed: \'$err\'')
        return ''
    }

    mut combined := []u8{len: nonce.len + ciphertext.len}
    combined << nonce
    combined << ciphertext
    data_b64 := base64.encode(combined)

    return data_b64
}

fn ascon_decrypt(key_hex string, data_b64 string) string {
    key := []u8(key_hex.hex())
    combined := base64.decode(data_b64)

    if combined.len < ascon.nonce_size {
        logs.error('cryptography:cryptography:ascon_decrypt', 'data blob is too short')
        return ''
    }

    nonce := combined[..ascon.nonce_size]
    ciphertext := combined[ascon.nonce_size..]

    plaintext_bytes := ascon.decrypt(key, nonce, []u8{}, ciphertext) or {
        logs.error('cryptography:cryptography:ascon_decrypt', 'decrypt failed: \'$err\'')
        return ''
    }

    return plaintext_bytes.bytestr()
}

fn generate_keypair() !structs.KeyPair {
    mut secret_key_bytes := curve25519.PrivateKey.new() or {
        logs.error('cryptography:cryptography:generate_keypair', 'x25519 secret key generation failed due to: \'$err\'')
        return error('x25519 secret key generation failed due to: \'$err\'')
    }
    public_key_bytes := secret_key_bytes.public_key() or {
        logs.error('cryptography:cryptography:generate_keypair', 'x25519 public key computation/generation failed due to: \'$err\'')
        return error('x25519 public key computation/generation failed due to: \'$err\'')
    }

    logs.info('cryptography:cryptography:generate_keypair', 'Generated x25519 keypair')
    return structs.KeyPair{
        secret: secret_key_bytes
        public: public_key_bytes
    }
}

fn compute_shared_secret(mut secret_key curve25519.PrivateKey, public_key curve25519.PublicKey) ![]u8 {
    shared_secret_raw := curve25519.derive_shared_secret(mut secret_key, public_key) or {
        logs.error('cryptography:cryptography:compute_shared_secret', 'x25519 raw shared_secret computation/derivation failed due to: \'$err\'')
        return []u8{}
    }

    logs.info('cryptography:cryptography:compute_shared_secret', 'Computed x25519 raw shared_secret')
    return shared_secret_raw
}

fn hkdf_sha256(ikm []u8, salt []u8, length int) []u8 {
    prk := hmac.new(salt, ikm, sha256.sum, sha256.block_size)

    info := []u8{}
    mut okm := []u8{}
    mut previous_block := []u8{}
    mut counter := u8(1)

    for okm.len < length {
        mut block_input := previous_block.clone()
        block_input << info
        block_input << counter

        block := hmac.new(prk, block_input, sha256.sum, sha256.block_size)

        okm << block
        previous_block = block.clone()
        counter++
    }

    logs.info('cryptography:cryptography:hkdf_sha256', 'Derived HMAC-x25519 shared_secret')
    return okm[..length]
}

fn salt_randomizer() []u8 {
    logs.info('cryptography:cryptography:salt_randomizer', 'Generated CSPRNG randomized salt')
    return rand.bytes(32) or {
        logs.error('cryptography:cryptography:salt_randomizer', 'CSPRNG randomized salt generation failed due to: \'$err\'')
        return []u8{}
    }
}