/*
 * LITLS -- ChaCha20-Poly1305 AEAD (RFC 8439)
 * Uses incremental Poly1305 to avoid VLA allocation.
 */

#include "litls.h"
#include <string.h>

static void pad16_update(litls_poly1305_ctx *ctx, size_t data_len) {
    size_t rem = data_len % 16;
    if (rem != 0) {
        uint8_t zeros[15];
        size_t pad = 16 - rem;
        memset(zeros, 0, pad);
        litls_poly1305_update(ctx, zeros, pad);
    }
}

static void le64_buf(uint8_t out[8], uint64_t v) {
    for (int i = 0; i < 8; i++)
        out[i] = (uint8_t)(v >> (i * 8));
}

int litls_chacha20_poly1305_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 pt_len, uint8_t *ciphertext, uint8_t tag[16]) {
    /* Generate Poly1305 one-time key */
    uint8_t poly_key[64];
    litls_chacha20_block(key, nonce, 0, poly_key);

    /* Encrypt plaintext (counter starts at 1) */
    litls_chacha20_encrypt(key, nonce, 1, plaintext, pt_len, ciphertext);

    /* Compute Poly1305 tag incrementally: aad || pad || ct || pad || len(aad) || len(ct) */
    litls_poly1305_ctx pctx;
    litls_poly1305_init(&pctx, poly_key);

    litls_poly1305_update(&pctx, aad, aad_len);
    pad16_update(&pctx, aad_len);

    litls_poly1305_update(&pctx, ciphertext, pt_len);
    pad16_update(&pctx, pt_len);

    uint8_t lens[16];
    le64_buf(lens, (uint64_t)aad_len);
    le64_buf(lens + 8, (uint64_t)pt_len);
    litls_poly1305_update(&pctx, lens, 16);

    litls_poly1305_final(&pctx, tag);
    return 0;
}

int litls_chacha20_poly1305_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 ct_len, const uint8_t tag[16],
                                    uint8_t *plaintext) {
    /* Generate Poly1305 one-time key */
    uint8_t poly_key[64];
    litls_chacha20_block(key, nonce, 0, poly_key);

    /* Verify tag first using incremental Poly1305 */
    litls_poly1305_ctx pctx;
    litls_poly1305_init(&pctx, poly_key);

    litls_poly1305_update(&pctx, aad, aad_len);
    pad16_update(&pctx, aad_len);

    litls_poly1305_update(&pctx, ciphertext, ct_len);
    pad16_update(&pctx, ct_len);

    uint8_t lens[16];
    le64_buf(lens, (uint64_t)aad_len);
    le64_buf(lens + 8, (uint64_t)ct_len);
    litls_poly1305_update(&pctx, lens, 16);

    uint8_t computed_tag[16];
    litls_poly1305_final(&pctx, computed_tag);

    if (litls_secure_memcmp(computed_tag, tag, 16) != 0)
        return -1;

    /* Decrypt */
    litls_chacha20_encrypt(key, nonce, 1, ciphertext, ct_len, plaintext);
    return 0;
}