#ifndef ZERIN_CRYPTO_H
#define ZERIN_CRYPTO_H

#include <stdint.h>
#include <stddef.h>

#define CRYPTO_KEY_SIZE     32
#define CRYPTO_NONCE_SIZE   12
#define CRYPTO_TAG_SIZE     16
#define CRYPTO_SHA256_SIZE  32
#define CRYPTO_X25519_KEY   32

// ============================================================================
// Random
// ============================================================================
int crypto_random_bytes(uint8_t *buf, size_t len);
double crypto_random_float(void);  // [0.0, 1.0)

// ============================================================================
// SHA-256
// ============================================================================
typedef struct {
    uint32_t state[8];
    uint64_t bitcount;
    uint8_t  buffer[64];
    uint32_t buflen;
} sha256_ctx_t;

void sha256_init(sha256_ctx_t *ctx);
void sha256_update(sha256_ctx_t *ctx, const uint8_t *data, size_t len);
void sha256_final(sha256_ctx_t *ctx, uint8_t out[32]);
void sha256(const uint8_t *data, size_t len, uint8_t out[32]);

// HMAC-SHA256
void hmac_sha256(const uint8_t *key, size_t key_len,
                 const uint8_t *data, size_t data_len,
                 uint8_t out[32]);

// ============================================================================
// HKDF (RFC 5869)
// ============================================================================
int hkdf_extract(const uint8_t *salt, size_t salt_len,
                 const uint8_t *ikm, size_t ikm_len,
                 uint8_t prk[32]);

int hkdf_expand(const uint8_t prk[32],
                const uint8_t *info, size_t info_len,
                uint8_t *okm, size_t okm_len);

int hkdf_derive(const uint8_t *salt, size_t salt_len,
                const uint8_t *ikm, size_t ikm_len,
                const uint8_t *info, size_t info_len,
                uint8_t *okm, size_t okm_len);

// ============================================================================
// X25519 (RFC 7748)
// ============================================================================
void x25519_clamp(uint8_t key[32]);
int  x25519_keygen(uint8_t privkey[32], uint8_t pubkey[32]);
int  x25519_shared_secret(const uint8_t privkey[32],
                          const uint8_t peer_pubkey[32],
                          uint8_t shared[32]);

// ============================================================================
// ChaCha20 (RFC 8439)
// ============================================================================
typedef struct {
    uint32_t state[16];
} chacha20_ctx_t;

void chacha20_init(chacha20_ctx_t *ctx, const uint8_t key[32],
                   const uint8_t nonce[12], uint32_t counter);
void chacha20_encrypt(chacha20_ctx_t *ctx, const uint8_t *in,
                      uint8_t *out, size_t len);

// ============================================================================
// Poly1305 (RFC 8439)
// ============================================================================
typedef struct {
    uint32_t r[5];
    uint32_t h[5];
    uint32_t pad[4];
    uint8_t  buffer[16];
    size_t   buflen;
    size_t   total;
} poly1305_ctx_t;

void poly1305_init(poly1305_ctx_t *ctx, const uint8_t key[32]);
void poly1305_update(poly1305_ctx_t *ctx, const uint8_t *data, size_t len);
void poly1305_final(poly1305_ctx_t *ctx, uint8_t tag[16]);

// ============================================================================
// ChaCha20-Poly1305 AEAD (RFC 8439)
// ============================================================================

// Encrypt: returns ciphertext_len (plaintext_len + 16 for tag)
// out must have room for plaintext_len + CRYPTO_TAG_SIZE
int aead_encrypt(const uint8_t key[32], const uint8_t nonce[12],
                 const uint8_t *aad, size_t aad_len,
                 const uint8_t *plaintext, size_t plaintext_len,
                 uint8_t *out);

// Decrypt: returns 0 on success, -1 on auth failure
// out must have room for ciphertext_len - CRYPTO_TAG_SIZE
int aead_decrypt(const uint8_t key[32], const uint8_t nonce[12],
                 const uint8_t *aad, size_t aad_len,
                 const uint8_t *ciphertext, size_t ciphertext_len,
                 uint8_t *out);

// ============================================================================
// Key Exchange
// ============================================================================
typedef struct {
    uint8_t agent_pubkey_hash[32];  // SHA256(agent_pubkey)
    uint8_t ephemeral_pubkey[32];
    uint8_t ephemeral_privkey[32];  // NOT sent
    uint8_t session_key[32];
} key_exchange_t;

int key_exchange_init(key_exchange_t *kx, const uint8_t agent_pubkey[32]);
int key_exchange_derive(key_exchange_t *kx, const uint8_t server_pubkey[32]);

// ============================================================================
// Utility
// ============================================================================
void secure_zero(void *ptr, size_t len);
int  constant_time_compare(const uint8_t *a, const uint8_t *b, size_t len);

#endif // ZERIN_CRYPTO_H