#include "curve25519.h"

#include <array>
#include <cstdint>
#include <vector>
#include <algorithm>

namespace netplus {

// ============================================================================
// FieldElement: gf[16] like TweetNaCl (internal only)
// ============================================================================

class FieldElement {
public:
    std::array<int64_t, 16> limbs;

    FieldElement() { limbs.fill(0); }
    explicit FieldElement(int v) { limbs.fill(0); limbs[0] = v; }

    void zero() { limbs.fill(0); }
    void one()  { limbs.fill(0); limbs[0] = 1; }

    // A(o,a,b): o=a+b   (NO carry, TweetNaCl style)
    FieldElement operator+(const FieldElement& other) const {
        FieldElement r;
        for (int i = 0; i < 16; ++i)
            r.limbs[i] = limbs[i] + other.limbs[i];
        return r;
    }

    // Z(o,a,b): o=a-b   (NO carry, TweetNaCl style)
    FieldElement operator-(const FieldElement& other) const {
        FieldElement r;
        for (int i = 0; i < 16; ++i)
            r.limbs[i] = limbs[i] - other.limbs[i];
        return r;
    }

    // car25519(o)  --- EXACT TweetNaCl
    void carry25519() {
        for (int i = 0; i < 16; ++i) {
            limbs[i] += (1LL << 16);
            int64_t c = limbs[i] >> 16;
            limbs[i] -= c << 16;

            if (i < 15) {
                limbs[i + 1] += c - 1;
            } else {
                limbs[0] += 38 * (c - 1);   // <-- WICHTIG: 38*(c-1), nicht 37*(c-1)
            }
        }
    }

    // sel25519(p,q,b) : swap if b=1  --- EXACT TweetNaCl
    void sel25519(FieldElement& q, int b) {
        int64_t t;
        int64_t c = ~((int64_t)b - 1); // b=0 => 0, b=1 => -1
        for (int i = 0; i < 16; ++i) {
            t = c & (limbs[i] ^ q.limbs[i]);
            limbs[i] ^= t;
            q.limbs[i] ^= t;
        }
    }

    // M(o,a,b): multiplication + reduction  --- EXACT TweetNaCl
    FieldElement operator*(const FieldElement& g) const {
        int64_t t[31];
        for (int i = 0; i < 31; ++i) t[i] = 0;

        for (int i = 0; i < 16; ++i)
            for (int j = 0; j < 16; ++j)
                t[i + j] += limbs[i] * g.limbs[j];

        for (int i = 0; i < 15; ++i)
            t[i] += 38 * t[i + 16];

        FieldElement out;
        for (int i = 0; i < 16; ++i)
            out.limbs[i] = t[i];

        out.carry25519();
        out.carry25519();
        return out;
    }

    // S(o,a): optimized square using symmetry (136 muls instead of 256)
    FieldElement square() const {
        int64_t t[31];
        for (int i = 0; i < 31; ++i) t[i] = 0;

        // Off-diagonal: t[i+j] += 2 * limbs[i] * limbs[j] for i < j
        for (int i = 0; i < 15; ++i)
            for (int j = i + 1; j < 16; ++j)
                t[i + j] += limbs[i] * limbs[j];

        // Double the off-diagonal contributions
        for (int i = 0; i < 31; ++i) t[i] *= 2;

        // Diagonal: t[2*i] += limbs[i]^2
        for (int i = 0; i < 16; ++i)
            t[2 * i] += limbs[i] * limbs[i];

        // Reduce and carry (same as operator*)
        for (int i = 0; i < 15; ++i)
            t[i] += 38 * t[i + 16];

        FieldElement out;
        for (int i = 0; i < 16; ++i)
            out.limbs[i] = t[i];

        out.carry25519();
        out.carry25519();
        return out;
    }

    // unpack25519
    void from_bytes(const std::vector<uint8_t>& in) {
        for (int i = 0; i < 16; ++i) {
            limbs[i] = (int64_t)in[2*i] + ((int64_t)in[2*i+1] << 8);
        }
        limbs[15] &= 0x7fff;
    }

    // pack25519  --- EXACT TweetNaCl
    std::vector<uint8_t> to_bytes() const {
        std::vector<uint8_t> out(32, 0);
        FieldElement t = *this;
        FieldElement m;

        t.carry25519();
        t.carry25519();
        t.carry25519();

        for (int j = 0; j < 2; ++j) {
            m.limbs[0] = t.limbs[0] - 0xffed;
            for (int i = 1; i < 15; ++i) {
                m.limbs[i] = t.limbs[i] - 0xffff - ((m.limbs[i-1] >> 16) & 1);
                m.limbs[i-1] &= 0xffff;
            }
            m.limbs[15] = t.limbs[15] - 0x7fff - ((m.limbs[14] >> 16) & 1);
            int b = (m.limbs[15] >> 16) & 1;
            m.limbs[14] &= 0xffff;
            t.sel25519(m, 1 - b);
        }

        for (int i = 0; i < 16; ++i) {
            out[2*i]     = (uint8_t)(t.limbs[i] & 0xff);
            out[2*i + 1] = (uint8_t)(t.limbs[i] >> 8);
        }
        return out;
    }

    // inv25519(o,i)  --- EXACT TweetNaCl
    FieldElement invert() const {
        FieldElement c = *this;
        for (int a = 253; a >= 0; --a) {
            c = c.square();
            if (a != 2 && a != 4)
                c = c * (*this);
        }
        return c;
    }
};

// ============================================================================
// X25519 scalar multiplication (TweetNaCl crypto_scalarmult)
// ============================================================================

bool scalarmult_curve25519(std::vector<uint8_t>& out,
                           const std::vector<uint8_t>& scalar,
                           const std::vector<uint8_t>& point_u)
{
    if (scalar.size() != 32 || point_u.size() != 32) return false;

    // Clamp scalar
    std::array<uint8_t, 32> z{};
    for (int i = 0; i < 32; ++i) z[i] = scalar[i];
    z[31] = (z[31] & 127) | 64;
    z[0] &= 248;


    // Decode u
    FieldElement x;
    x.from_bytes(point_u);

    // Ladder state
    FieldElement a, b, c, d, e, f;
    b = x;
    a.zero();
    c.zero();
    d.zero();
    a.limbs[0] = 1;
    d.limbs[0] = 1;

    FieldElement _121665;
    _121665.zero();
    _121665.limbs[0] = 0xDB41;
    _121665.limbs[1] = 1;

    // Main ladder loop (exact TweetNaCl)
    for (int i = 254; i >= 0; --i) {
        int r = (z[i >> 3] >> (i & 7)) & 1;

        a.sel25519(b, r);
        c.sel25519(d, r);

        e = a + c;
        a = a - c;
        c = b + d;
        b = b - d;

        d = e.square();
        f = a.square();

        a = c * a;
        c = b * e;

        e = a + c;
        a = a - c;

        b = a.square();
        c = d - f;

        a = c * _121665;
        a = a + d;

        c = c * a;
        a = d * f;

        d = b * x;
        b = e.square();

        a.sel25519(b, r);
        c.sel25519(d, r);
    }

    // Convert: a / c
    c = c.invert();
    a = a * c;

    out = a.to_bytes();
    return true;
}

bool scalarmult_curve25519_base(std::vector<uint8_t>& out,
                                const std::vector<uint8_t>& scalar)
{
    if (scalar.size() != 32) return false;

    std::vector<uint8_t> base(32, 0);
    base[0] = 9;
    return scalarmult_curve25519(out, scalar, base);
}

} // namespace netplus
