/*
 * LITLS -- Ed25519 sign/verify
 * Extracted from TweetNaCl (public domain, DJB / Bernstein, Schwabe, et al.)
 * Uses litls_sha512 instead of TweetNaCl's internal SHA-512.
 * Single-letter variable names expanded for readability.
 */

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

typedef int64_t fe[16]; /* Field element in GF(2^255-19) */

/* ── GF(2^255-19) field arithmetic (shared with X25519) ────────── */

static const fe fe_zero = {0};
static const fe fe_one  = {1};

/* Ed25519 curve constant d = -121665/121666 mod p */
static const fe ed_d = {0x78a3, 0x1359, 0x4dca, 0x75eb, 0xd8ab, 0x4141, 0x0a4d, 0x0070,
                        0xe898, 0x7779, 0x4079, 0x8cc7, 0xfe73, 0x2b6f, 0x6cee, 0x5203};

/* 2*d */
static const fe ed_d2 = {0xf159, 0x26b2, 0x9b94, 0xebd6, 0xb156, 0x8283, 0x149a, 0x00e0,
                         0xd130, 0xeef3, 0x80f2, 0x198e, 0xfce7, 0x56df, 0xd9dc, 0x2406};

/* Base point x-coordinate */
static const fe ed_base_x = {0xd51a, 0x8f25, 0x2d60, 0xc956, 0xa7b2, 0x9525, 0xc760, 0x692c,
                             0xdc5c, 0xfdd6, 0xe231, 0xc0a4, 0x53fe, 0xcd6e, 0x36d3, 0x2169};

/* Base point y-coordinate */
static const fe ed_base_y = {0x6658, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666,
                             0x6666, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666, 0x6666};

/* sqrt(-1) mod p */
static const fe fe_sqrtm1 = {0xa0b0, 0x4a0e, 0x1b27, 0xc4ee, 0xe478, 0xad2f, 0x1806, 0x2f43,
                             0xd7a7, 0x3dfb, 0x0099, 0x2b4d, 0xdf0b, 0x4fc1, 0x2480, 0x2b83};

/* Group order L = 2^252 + 27742317777372353535851937790883648493 */
static const uint64_t sc_L[32] = {0xed, 0xd3, 0xf5, 0x5c, 0x1a, 0x63, 0x12, 0x58, 0xd6, 0x9c, 0xf7,
                                  0xa2, 0xde, 0xf9, 0xde, 0x14, 0,    0,    0,    0,    0,    0,
                                  0,    0,    0,    0,    0,    0,    0,    0,    0,    0x10};

static void fe_copy(fe out, const fe a) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i];
}

static void fe_carry(fe o) {
    int i;
    int64_t c;
    for (i = 0; i < 16; i++) {
        o[i] += (1LL << 16);
        c = o[i] >> 16;
        o[(i + 1) * (i < 15)] += c - 1 + 37 * (c - 1) * (i == 15);
        o[i] -= c << 16;
    }
}

static void fe_cswap(fe p, fe q, int b) {
    int64_t t, c = ~(b - 1);
    int i;
    for (i = 0; i < 16; i++) {
        t = c & (p[i] ^ q[i]);
        p[i] ^= t;
        q[i] ^= t;
    }
}

static void fe_pack(uint8_t out[32], const fe n) {
    int i, j, b;
    fe m, t;
    for (i = 0; i < 16; i++)
        t[i] = n[i];
    fe_carry(t);
    fe_carry(t);
    fe_carry(t);
    for (j = 0; j < 2; j++) {
        m[0] = t[0] - 0xffed;
        for (i = 1; i < 15; i++) {
            m[i] = t[i] - 0xffff - ((m[i - 1] >> 16) & 1);
            m[i - 1] &= 0xffff;
        }
        m[15] = t[15] - 0x7fff - ((m[14] >> 16) & 1);
        b     = (m[15] >> 16) & 1;
        m[14] &= 0xffff;
        fe_cswap(t, m, 1 - b);
    }
    for (i = 0; i < 16; i++) {
        out[2 * i]     = t[i] & 0xff;
        out[2 * i + 1] = t[i] >> 8;
    }
}

static void fe_unpack(fe out, const uint8_t n[32]) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = n[2 * i] + ((int64_t)n[2 * i + 1] << 8);
    out[15] &= 0x7fff;
}

static void fe_add(fe out, const fe a, const fe b) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i] + b[i];
}

static void fe_sub(fe out, const fe a, const fe b) {
    int i;
    for (i = 0; i < 16; i++)
        out[i] = a[i] - b[i];
}

static void fe_mul(fe out, const fe a, const fe b) {
    int64_t i, j, t[31];
    for (i = 0; i < 31; i++)
        t[i] = 0;
    for (i = 0; i < 16; i++)
        for (j = 0; j < 16; j++)
            t[i + j] += a[i] * b[j];
    for (i = 0; i < 15; i++)
        t[i] += 38 * t[i + 16];
    for (i = 0; i < 16; i++)
        out[i] = t[i];
    fe_carry(out);
    fe_carry(out);
}

static void fe_sq(fe out, const fe a) {
    fe_mul(out, a, a);
}

static void fe_inv(fe out, const fe in) {
    fe c;
    int a;
    for (a = 0; a < 16; a++)
        c[a] = in[a];
    for (a = 253; a >= 0; a--) {
        fe_sq(c, c);
        if (a != 2 && a != 4)
            fe_mul(c, c, in);
    }
    for (a = 0; a < 16; a++)
        out[a] = c[a];
}

/* ── Ed25519-specific field ops ──────────────────────────── */

static int fe_neq(const fe a, const fe b) {
    uint8_t c[32], d[32];
    fe_pack(c, a);
    fe_pack(d, b);
    return litls_secure_memcmp(c, d, 32);
}

static uint8_t fe_parity(const fe a) {
    uint8_t d[32];
    fe_pack(d, a);
    return d[0] & 1;
}

/* Compute a^(2^252 - 3) — used for square root recovery */
static void fe_pow2523(fe out, const fe in) {
    fe c;
    int a;
    for (a = 0; a < 16; a++)
        c[a] = in[a];
    for (a = 250; a >= 0; a--) {
        fe_sq(c, c);
        if (a != 1)
            fe_mul(c, c, in);
    }
    for (a = 0; a < 16; a++)
        out[a] = c[a];
}

/* ── Extended point operations (X, Y, Z, T) ──────────────── */

/* Point addition in extended coordinates */
static void ge_add(fe p[4], fe q[4]) {
    fe a, b, c, d, t, e, f, g, h;

    fe_sub(a, p[1], p[0]);
    fe_sub(t, q[1], q[0]);
    fe_mul(a, a, t);
    fe_add(b, p[0], p[1]);
    fe_add(t, q[0], q[1]);
    fe_mul(b, b, t);
    fe_mul(c, p[3], q[3]);
    fe_mul(c, c, ed_d2);
    fe_mul(d, p[2], q[2]);
    fe_add(d, d, d);
    fe_sub(e, b, a);
    fe_sub(f, d, c);
    fe_add(g, d, c);
    fe_add(h, b, a);

    fe_mul(p[0], e, f);
    fe_mul(p[1], h, g);
    fe_mul(p[2], g, f);
    fe_mul(p[3], e, h);
}

/* Constant-time point swap */
static void ge_cswap(fe p[4], fe q[4], uint8_t b) {
    int i;
    for (i = 0; i < 4; i++)
        fe_cswap(p[i], q[i], b);
}

/* Pack extended point to 32 bytes (compress to y with sign bit) */
static void ge_pack(uint8_t out[32], fe p[4]) {
    fe tx, ty, zi;
    fe_inv(zi, p[2]);
    fe_mul(tx, p[0], zi);
    fe_mul(ty, p[1], zi);
    fe_pack(out, ty);
    out[31] ^= fe_parity(tx) << 7;
}

/* Scalar multiplication: p = s * q */
static void ge_scalarmult(fe p[4], fe q[4], const uint8_t *s) {
    int i;
    fe_copy(p[0], fe_zero);
    fe_copy(p[1], fe_one);
    fe_copy(p[2], fe_one);
    fe_copy(p[3], fe_zero);
    for (i = 255; i >= 0; --i) {
        uint8_t b = (s[i / 8] >> (i & 7)) & 1;
        ge_cswap(p, q, b);
        ge_add(q, p);
        ge_add(p, p);
        ge_cswap(p, q, b);
    }
}

/* Scalar multiplication by the base point */
static void ge_scalarbase(fe p[4], const uint8_t *s) {
    fe q[4];
    fe_copy(q[0], ed_base_x);
    fe_copy(q[1], ed_base_y);
    fe_copy(q[2], fe_one);
    fe_mul(q[3], ed_base_x, ed_base_y);
    ge_scalarmult(p, q, s);
}

/* Unpack a compressed point and negate it (for verification) */
static int ge_unpack_neg(fe r[4], const uint8_t p[32]) {
    fe t, chk, num, den, den2, den4, den6;
    fe_copy(r[2], fe_one);
    fe_unpack(r[1], p);
    fe_sq(num, r[1]);
    fe_mul(den, num, ed_d);
    fe_sub(num, num, r[2]);
    fe_add(den, r[2], den);

    fe_sq(den2, den);
    fe_sq(den4, den2);
    fe_mul(den6, den4, den2);
    fe_mul(t, den6, num);
    fe_mul(t, t, den);

    fe_pow2523(t, t);
    fe_mul(t, t, num);
    fe_mul(t, t, den);
    fe_mul(t, t, den);
    fe_mul(r[0], t, den);

    fe_sq(chk, r[0]);
    fe_mul(chk, chk, den);
    if (fe_neq(chk, num))
        fe_mul(r[0], r[0], fe_sqrtm1);

    fe_sq(chk, r[0]);
    fe_mul(chk, chk, den);
    if (fe_neq(chk, num))
        return -1;

    if (fe_parity(r[0]) == (p[31] >> 7))
        fe_sub(r[0], fe_zero, r[0]);

    fe_mul(r[3], r[0], r[1]);
    return 0;
}

/* ── Scalar arithmetic mod L ──────────────────────────────── */

/* Reduce x[64] mod L, output to r[32] */
static void sc_modL(uint8_t *r, int64_t x[64]) {
    int64_t carry, i, j;
    for (i = 63; i >= 32; --i) {
        carry = 0;
        for (j = i - 32; j < i - 12; ++j) {
            x[j] += carry - 16 * x[i] * (int64_t)sc_L[j - (i - 32)];
            carry = (x[j] + 128) >> 8;
            x[j] -= carry << 8;
        }
        x[j] += carry;
        x[i] = 0;
    }
    carry = 0;
    for (j = 0; j < 32; j++) {
        x[j] += carry - (x[31] >> 4) * (int64_t)sc_L[j];
        carry = x[j] >> 8;
        x[j] &= 255;
    }
    for (j = 0; j < 32; j++)
        x[j] -= carry * (int64_t)sc_L[j];
    for (i = 0; i < 32; i++) {
        x[i + 1] += x[i] >> 8;
        r[i] = x[i] & 255;
    }
}

/* Reduce a 64-byte hash mod L in place */
static void sc_reduce(uint8_t r[64]) {
    int64_t x[64];
    int i;
    for (i = 0; i < 64; i++)
        x[i] = (uint64_t)r[i];
    for (i = 0; i < 64; i++)
        r[i] = 0;
    sc_modL(r, x);
}

/* ── Public API ───────────────────────────────────────────── */

void litls_ed25519_keypair(uint8_t public_key[32], uint8_t secret_key[64], const uint8_t seed[32]) {
    uint8_t d[64];
    fe p[4];

    litls_sha512(seed, 32, d);
    d[0] &= 248;
    d[31] &= 127;
    d[31] |= 64;

    ge_scalarbase(p, d);
    ge_pack(public_key, p);

    memcpy(secret_key, seed, 32);
    memcpy(secret_key + 32, public_key, 32);
}

int litls_ed25519_sign(uint8_t sig[64], const uint8_t *msg, size_t msg_len, const uint8_t public_key[32],
                       const uint8_t secret_key[64]) {
    uint8_t d[64], nonce[64], hram[64];
    int64_t x[64];
    int i, j;
    fe p[4];
    litls_sha512_ctx ctx;

    /* Expand secret key (seed is first 32 bytes) */
    litls_sha512(secret_key, 32, d);
    d[0] &= 248;
    d[31] &= 127;
    d[31] |= 64;

    /* nonce = SHA-512(d[32..63] || msg) */
    litls_sha512_init(&ctx);
    litls_sha512_update(&ctx, d + 32, 32);
    litls_sha512_update(&ctx, msg, msg_len);
    litls_sha512_final(&ctx, nonce);
    sc_reduce(nonce);

    /* R = nonce * B */
    ge_scalarbase(p, nonce);
    ge_pack(sig, p);

    /* hram = SHA-512(R || public_key || msg) */
    litls_sha512_init(&ctx);
    litls_sha512_update(&ctx, sig, 32);
    litls_sha512_update(&ctx, public_key, 32);
    litls_sha512_update(&ctx, msg, msg_len);
    litls_sha512_final(&ctx, hram);
    sc_reduce(hram);

    /* S = (nonce + hram * d) mod L */
    for (i = 0; i < 64; i++)
        x[i] = 0;
    for (i = 0; i < 32; i++)
        x[i] = (uint64_t)nonce[i];
    for (i = 0; i < 32; i++)
        for (j = 0; j < 32; j++)
            x[i + j] += hram[i] * (uint64_t)d[j];
    sc_modL(sig + 32, x);

    return 0;
}

/* ── Ed25519 to X25519 conversion ────────────────────────── */

int litls_ed25519_to_x25519_pk(uint8_t x25519_pk[32], const uint8_t ed25519_pk[32]) {
    fe y, one_plus_y, one_minus_y, inv, u;

    /* Unpack y from compressed Ed25519 point (ignore sign bit) */
    fe_unpack(y, ed25519_pk);

    /* u = (1 + y) / (1 - y) mod p */
    fe_add(one_plus_y, fe_one, y);
    fe_sub(one_minus_y, fe_one, y);
    fe_inv(inv, one_minus_y);
    fe_mul(u, one_plus_y, inv);

    fe_pack(x25519_pk, u);
    return 0;
}

int litls_ed25519_to_x25519_sk(uint8_t x25519_sk[32], const uint8_t ed25519_seed[32]) {
    uint8_t h[64];

    /* Same expansion as Ed25519 keypair: SHA-512(seed), clamp, take first 32 */
    litls_sha512(ed25519_seed, 32, h);
    h[0] &= 248;
    h[31] &= 127;
    h[31] |= 64;

    memcpy(x25519_sk, h, 32);
    litls_secure_zero(h, 64);
    return 0;
}

int litls_ed25519_verify(const uint8_t sig[64], const uint8_t *msg, size_t msg_len, const uint8_t public_key[32]) {
    uint8_t t[32], h[64];
    fe p[4], q[4];
    litls_sha512_ctx ctx;

    if (ge_unpack_neg(q, public_key))
        return -1;

    /* h = SHA-512(R || public_key || msg) */
    litls_sha512_init(&ctx);
    litls_sha512_update(&ctx, sig, 32);
    litls_sha512_update(&ctx, public_key, 32);
    litls_sha512_update(&ctx, msg, msg_len);
    litls_sha512_final(&ctx, h);
    sc_reduce(h);

    /* Check: S*B == R + h*A
     * q = -A (from unpack_neg), so compute h*(-A) + S*B and compare to R */
    ge_scalarmult(p, q, h);
    ge_scalarbase(q, sig + 32);
    ge_add(p, q);
    ge_pack(t, p);

    if (litls_secure_memcmp(sig, t, 32) != 0)
        return -1;

    return 0;
}