/*
 * LITLS -- Poly1305 MAC (RFC 8439)
 * Based on DJB's/Andrew Moon's approach, using 64-bit arithmetic for accumulator.
 */

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

static uint64_t le_load32(const uint8_t *p) {
    return (uint64_t)p[0] | ((uint64_t)p[1] << 8) | ((uint64_t)p[2] << 16) | ((uint64_t)p[3] << 24);
}

/* Process one full 16-byte block (with hibit set) */
static void poly1305_block(litls_poly1305_ctx *ctx, const uint8_t block[16], int final_partial, size_t blen) {
    uint8_t padded[17];
    memcpy(padded, block, blen);
    if (!final_partial) {
        /* Full block: hibit at position 16 */
        (void)blen;
        blen = 16;
    }
    padded[blen] = 1;
    memset(padded + blen + 1, 0, 16 - blen);

    uint32_t t0 = le_load32(padded + 0) & 0x03ffffff;
    uint32_t t1 = (le_load32(padded + 3) >> 2) & 0x03ffffff;
    uint32_t t2 = (le_load32(padded + 6) >> 4) & 0x03ffffff;
    uint32_t t3 = (le_load32(padded + 9) >> 6) & 0x03ffffff;
    uint32_t t4 = (le_load32(padded + 12) >> 8);
    if (!final_partial)
        t4 |= (1 << 24);

    ctx->h[0] += t0;
    ctx->h[1] += t1;
    ctx->h[2] += t2;
    ctx->h[3] += t3;
    ctx->h[4] += t4;

    uint64_t d0 = (uint64_t)ctx->h[0] * ctx->r[0] + (uint64_t)ctx->h[1] * ctx->s5[3] +
                  (uint64_t)ctx->h[2] * ctx->s5[2] + (uint64_t)ctx->h[3] * ctx->s5[1] +
                  (uint64_t)ctx->h[4] * ctx->s5[0];
    uint64_t d1 = (uint64_t)ctx->h[0] * ctx->r[1] + (uint64_t)ctx->h[1] * ctx->r[0] + (uint64_t)ctx->h[2] * ctx->s5[3] +
                  (uint64_t)ctx->h[3] * ctx->s5[2] + (uint64_t)ctx->h[4] * ctx->s5[1];
    uint64_t d2 = (uint64_t)ctx->h[0] * ctx->r[2] + (uint64_t)ctx->h[1] * ctx->r[1] + (uint64_t)ctx->h[2] * ctx->r[0] +
                  (uint64_t)ctx->h[3] * ctx->s5[3] + (uint64_t)ctx->h[4] * ctx->s5[2];
    uint64_t d3 = (uint64_t)ctx->h[0] * ctx->r[3] + (uint64_t)ctx->h[1] * ctx->r[2] + (uint64_t)ctx->h[2] * ctx->r[1] +
                  (uint64_t)ctx->h[3] * ctx->r[0] + (uint64_t)ctx->h[4] * ctx->s5[3];
    uint64_t d4 = (uint64_t)ctx->h[0] * ctx->r[4] + (uint64_t)ctx->h[1] * ctx->r[3] + (uint64_t)ctx->h[2] * ctx->r[2] +
                  (uint64_t)ctx->h[3] * ctx->r[1] + (uint64_t)ctx->h[4] * ctx->r[0];

    uint32_t c;
    c         = (uint32_t)(d0 >> 26);
    ctx->h[0] = (uint32_t)d0 & 0x03ffffff;
    d1 += c;
    c         = (uint32_t)(d1 >> 26);
    ctx->h[1] = (uint32_t)d1 & 0x03ffffff;
    d2 += c;
    c         = (uint32_t)(d2 >> 26);
    ctx->h[2] = (uint32_t)d2 & 0x03ffffff;
    d3 += c;
    c         = (uint32_t)(d3 >> 26);
    ctx->h[3] = (uint32_t)d3 & 0x03ffffff;
    d4 += c;
    c         = (uint32_t)(d4 >> 26);
    ctx->h[4] = (uint32_t)d4 & 0x03ffffff;
    ctx->h[0] += c * 5;
    c = ctx->h[0] >> 26;
    ctx->h[0] &= 0x03ffffff;
    ctx->h[1] += c;
}

void litls_poly1305_init(litls_poly1305_ctx *ctx, const uint8_t key[32]) {
    /* Clamp r */
    uint32_t r0 = (uint32_t)(le_load32(key + 0)) & 0x0fffffff;
    uint32_t r1 = (uint32_t)(le_load32(key + 4)) & 0x0ffffffc;
    uint32_t r2 = (uint32_t)(le_load32(key + 8)) & 0x0ffffffc;
    uint32_t r3 = (uint32_t)(le_load32(key + 12)) & 0x0ffffffc;

    ctx->r[0] = r0 & 0x03ffffff;
    ctx->r[1] = ((r0 >> 26) | (r1 << 6)) & 0x03ffffff;
    ctx->r[2] = ((r1 >> 20) | (r2 << 12)) & 0x03ffffff;
    ctx->r[3] = ((r2 >> 14) | (r3 << 18)) & 0x03ffffff;
    ctx->r[4] = (r3 >> 8) & 0x03ffffff;

    ctx->s5[0] = ctx->r[1] * 5;
    ctx->s5[1] = ctx->r[2] * 5;
    ctx->s5[2] = ctx->r[3] * 5;
    ctx->s5[3] = ctx->r[4] * 5;

    ctx->pad[0] = (uint32_t)(le_load32(key + 16));
    ctx->pad[1] = (uint32_t)(le_load32(key + 20));
    ctx->pad[2] = (uint32_t)(le_load32(key + 24));
    ctx->pad[3] = (uint32_t)(le_load32(key + 28));

    ctx->h[0] = ctx->h[1] = ctx->h[2] = ctx->h[3] = ctx->h[4] = 0;
    ctx->buf_len                                              = 0;
}

void litls_poly1305_update(litls_poly1305_ctx *ctx, const uint8_t *data, size_t len) {
    /* Fill partial block first */
    if (ctx->buf_len > 0) {
        size_t need = 16 - ctx->buf_len;
        if (len < need) {
            memcpy(ctx->buf + ctx->buf_len, data, len);
            ctx->buf_len += len;
            return;
        }
        memcpy(ctx->buf + ctx->buf_len, data, need);
        poly1305_block(ctx, ctx->buf, 0, 16);
        ctx->buf_len = 0;
        data += need;
        len -= need;
    }

    /* Process full blocks */
    while (len >= 16) {
        poly1305_block(ctx, data, 0, 16);
        data += 16;
        len -= 16;
    }

    /* Buffer remainder */
    if (len > 0) {
        memcpy(ctx->buf, data, len);
        ctx->buf_len = len;
    }
}

void litls_poly1305_final(litls_poly1305_ctx *ctx, uint8_t tag[16]) {
    /* Process final partial block if any */
    if (ctx->buf_len > 0) {
        poly1305_block(ctx, ctx->buf, 1, ctx->buf_len);
    }

    /* Final reduction mod 2^130 - 5 */
    uint32_t h0 = ctx->h[0], h1 = ctx->h[1], h2 = ctx->h[2], h3 = ctx->h[3], h4 = ctx->h[4];

    uint32_t c = h1 >> 26;
    h1 &= 0x03ffffff;
    h2 += c;
    c = h2 >> 26;
    h2 &= 0x03ffffff;
    h3 += c;
    c = h3 >> 26;
    h3 &= 0x03ffffff;
    h4 += c;
    c = h4 >> 26;
    h4 &= 0x03ffffff;
    h0 += c * 5;
    c = h0 >> 26;
    h0 &= 0x03ffffff;
    h1 += c;

    /* Compute h + -p = h - (2^130 - 5) */
    uint32_t g0 = h0 + 5;
    c           = g0 >> 26;
    g0 &= 0x03ffffff;
    uint32_t g1 = h1 + c;
    c           = g1 >> 26;
    g1 &= 0x03ffffff;
    uint32_t g2 = h2 + c;
    c           = g2 >> 26;
    g2 &= 0x03ffffff;
    uint32_t g3 = h3 + c;
    c           = g3 >> 26;
    g3 &= 0x03ffffff;
    uint32_t g4 = h4 + c - (1 << 26);

    uint32_t mask = (g4 >> 31) - 1;
    h0            = (h0 & ~mask) | (g0 & mask);
    h1            = (h1 & ~mask) | (g1 & mask);
    h2            = (h2 & ~mask) | (g2 & mask);
    h3            = (h3 & ~mask) | (g3 & mask);
    h4            = (h4 & ~mask) | (g4 & mask);

    /* Assemble h from 26-bit limbs into 32-bit words */
    uint32_t w0 = h0 | (h1 << 26);
    uint32_t w1 = (h1 >> 6) | (h2 << 20);
    uint32_t w2 = (h2 >> 12) | (h3 << 14);
    uint32_t w3 = (h3 >> 18) | (h4 << 8);

    /* Add s with carry */
    uint64_t f0 = (uint64_t)w0 + ctx->pad[0];
    uint64_t f1 = (uint64_t)w1 + ctx->pad[1] + (f0 >> 32);
    uint64_t f2 = (uint64_t)w2 + ctx->pad[2] + (f1 >> 32);
    uint64_t f3 = (uint64_t)w3 + ctx->pad[3] + (f2 >> 32);

    /* Write tag in little-endian */
    tag[0]  = (uint8_t)(f0);
    tag[1]  = (uint8_t)(f0 >> 8);
    tag[2]  = (uint8_t)(f0 >> 16);
    tag[3]  = (uint8_t)(f0 >> 24);
    tag[4]  = (uint8_t)(f1);
    tag[5]  = (uint8_t)(f1 >> 8);
    tag[6]  = (uint8_t)(f1 >> 16);
    tag[7]  = (uint8_t)(f1 >> 24);
    tag[8]  = (uint8_t)(f2);
    tag[9]  = (uint8_t)(f2 >> 8);
    tag[10] = (uint8_t)(f2 >> 16);
    tag[11] = (uint8_t)(f2 >> 24);
    tag[12] = (uint8_t)(f3);
    tag[13] = (uint8_t)(f3 >> 8);
    tag[14] = (uint8_t)(f3 >> 16);
    tag[15] = (uint8_t)(f3 >> 24);
}

/* One-shot API implemented via streaming */
void litls_poly1305_auth(const uint8_t key[32], const uint8_t *msg, size_t len, uint8_t tag[16]) {
    litls_poly1305_ctx ctx;
    litls_poly1305_init(&ctx, key);
    litls_poly1305_update(&ctx, msg, len);
    litls_poly1305_final(&ctx, tag);
}