#include "crypto/ecc_p256.h"
#include "crypto/ecc_u256.h"
#include <vector>
#include <cstdint>

namespace netplus {

// ------------------------------------------------------------
// Curve constants (NIST P-256 / secp256r1)
// ------------------------------------------------------------

// b coefficient (NORMAL domain literal, then converted once into Montgomery)

static const uint8_t P256_B_BE[32] = {
  0x5a,0xc6,0x35,0xd8,0xaa,0x3a,0x93,0xe7,
  0xb3,0xeb,0xbd,0x55,0x76,0x98,0x86,0xbc,
  0x65,0x1d,0x06,0xb0,0xcc,0x53,0xb0,0xf6,
  0x3b,0xce,0x3c,0x3e,0x27,0xd2,0x60,0x4b
};

const u256 P256_B = bytes_be_to_u256(P256_B_BE);

// NIST P-256 / secp256r1 base point G (uncompressed)
// Gx = 6b17d1f2e12c4247f8bce6e563a440f277037d812deb33a0f4a13945d898c296
// Gy = 4fe342e2fe1a7f9b8ee7eb4a7c0f9e162bce33576b315ececbb6406837bf51f5

// Standard basepoint (normal domain, LE u256)
static const P256Point P256_G_STD = []{
    P256Point g;
    g.inf  = false;
    g.mont = false;
    g.x = bytes_be_to_u256(P256_GX_BE);
    g.y = bytes_be_to_u256(P256_GY_BE);
    return g;
}();

const P256Point& P256_G() { return P256_G_STD; }


// Montgomery basepoint (for internal jacobian arithmetic)
static const P256Point P256_G_JAC_MONT = []{
    P256Point j;
    j.mont=true;
    j.inf = false;
    j.x = to_mont(P256_G_STD.x);
    j.y = to_mont(P256_G_STD.y);

    u256 one; u256_zero(one); one.w[0]=1;
    j.z = to_mont(one); // Z=1 in mont

    return j;
}();

// ------------------------------------------------------------
// Point Transformations
// ------------------------------------------------------------

/**
 * Converts Affine (Montgomery) to Jacobian (Montgomery).
 * Z starts as 1 in Montgomery domain: Z = to_mont(1).
 */
static inline P256Point jac_inf() {
    P256Point r;
    r.inf = true;
    r.mont = true;
    u256 zero; u256_zero(zero);
    r.x = to_mont(zero);
    r.y = to_mont(zero);
    r.z = to_mont(zero);  // Z=0 => infinity
    return r;
}

/**
 * Converts Jacobian (Montgomery) to Affine (Montgomery).
 * x = X / Z^2, y = Y / Z^3  (all Montgomery domain)
 */
P256Point toAff(const P256Point& Jm)
{
    P256Point P;
    P.mont = false;  // ✅ because we return STD

    if (Jm.inf) {
        P.inf = true;
        return P;
    }

    // Z == 0 ?
    u256 z_std = from_mont(Jm.z);
    if (u256_is_zero(z_std)) {
        P.inf = true;
        return P;
    }

    // zinv = 1/Z in MONT
    u256 zinv  = fp_inv(Jm.z);
    u256 zinv2 = fp_mul(zinv, zinv);
    u256 zinv3 = fp_mul(zinv2, zinv);

    // affine in MONT
    u256 Xm = fp_mul(Jm.x, zinv2);
    u256 Ym = fp_mul(Jm.y, zinv3);

    // convert to STD for output
    P.inf = false;
    P.x   = from_mont(Xm);
    P.y   = from_mont(Ym);
    return P;
}


// ------------------------------------------------------------
// Point Arithmetic
// ------------------------------------------------------------

/**
 * Jacobian point doubling: R = 2P  (all Montgomery domain)
 * Optimized for a=-3.
 */
P256Point point_double(const P256Point& a) {
    if (a.inf || u256_is_zero(a.z)) return a;

    u256 t1 = fp_sqr(a.z);              // Z^2
    u256 t2 = fp_sub(a.x, t1);          // X - Z^2
    u256 t3 = fp_add(a.x, t1);          // X + Z^2
    t2 = fp_mul(t2, t3);                // (X-Z^2)(X+Z^2)
    u256 t4 = fp_add(t2, t2);           // 2*t2
    u256 M  = fp_add(t4, t2);           // 3*(X^2 - Z^4)

    u256 Y2 = fp_sqr(a.y);              // Y^2
    u256 S  = fp_mul(a.x, Y2);          // X*Y^2
    S = fp_add(S, S);
    S = fp_add(S, S);                   // S = 4*X*Y^2

    u256 T  = fp_sqr(M);                // M^2
    u256 S2 = fp_add(S, S);             // 2S
    u256 X3 = fp_sub(T, S2);

    u256 Y4 = fp_sqr(Y2);
    u256 Y8 = fp_add(Y4, Y4);
    Y8 = fp_add(Y8, Y8);
    Y8 = fp_add(Y8, Y8);                // 8*Y^4

    u256 S_minus_X3 = fp_sub(S, X3);
    u256 Y3 = fp_sub(fp_mul(M, S_minus_X3), Y8);

    u256 YZ = fp_mul(a.y, a.z);
    u256 Z3 = fp_add(YZ, YZ);

    P256Point r;
    r.inf = false;
    r.x = X3;
    r.y = Y3;
    r.z = Z3;
    return r;
}

/**
 * Full Jacobian point addition: R = a + b (Montgomery)
 */
P256Point point_add_jac(const P256Point& a, const P256Point& b) {
    if (a.inf) return b;
    if (b.inf) return a;

    u256 z1z1 = fp_sqr(a.z);
    u256 z2z2 = fp_sqr(b.z);

    u256 u1 = fp_mul(a.x, z2z2);
    u256 u2 = fp_mul(b.x, z1z1);

    u256 s1 = fp_mul(fp_mul(a.y, b.z), z2z2);
    u256 s2 = fp_mul(fp_mul(b.y, a.z), z1z1);

    if (u256_cmp(u1, u2) == 0) {
        if (u256_cmp(s1, s2) == 0) return point_double(a);
        P256Point r; r.inf = true; return r;
    }

    u256 h = fp_sub(u2, u1);
    u256 r_val = fp_sub(s2, s1);
    u256 h2 = fp_sqr(h);
    u256 h3 = fp_mul(h2, h);
    u256 v  = fp_mul(u1, h2);

    u256 x3 = fp_sub(fp_sub(fp_sqr(r_val), h3), fp_add(v, v));
    u256 y3 = fp_sub(fp_mul(r_val, fp_sub(v, x3)), fp_mul(s1, h3));
    u256 z3 = fp_mul(fp_mul(a.z, b.z), h);

    P256Point res;
    res.inf = false;
    res.x = x3;
    res.y = y3;
    res.z = z3;
    return res;
}

// ------------------------------------------------------------
// Scalar Multiplication (Montgomery Ladder)
// ------------------------------------------------------------

static inline void cswap_jac(P256Point& A, P256Point& B, uint32_t swap) {
    // swap is 0 or 1
    cswap_u256(A.x, B.x, swap);
    cswap_u256(A.y, B.y, swap);
    cswap_u256(A.z, B.z, swap);

    // IMPORTANT: swap the infinity flag too
    uint32_t t = (uint32_t)A.inf ^ (uint32_t)B.inf;
    t &= (0u - swap);
    A.inf = (bool)((uint32_t)A.inf ^ t);
    B.inf = (bool)((uint32_t)B.inf ^ t);
}

P256Point toJac(const P256Point& Pstd)
{
    P256Point J;
    J.mont = true;

    if (Pstd.inf) {
        J.inf = true;
        u256 z; u256_zero(z);
        J.x = to_mont(z);
        J.y = to_mont(z);
        J.z = to_mont(z); // Z=0 => INF
        return J;
    }

    J.inf = false;
    J.x = to_mont(Pstd.x);
    J.y = to_mont(Pstd.y);

    u256 one; u256_zero(one); one.w[0]=1;
    J.z = to_mont(one); // Z=1
    return J;
}


static bool scalar_is_zero(const uint8_t k[32]) {
    for (int i=0;i<32;i++) if (k[i]) return false;
    return true;
}

P256Point scalar_mul(const uint8_t k[32], const P256Point& Pstd)
{
    // R = INF (Jacobian, MONT)
    P256Point R;
    R.inf  = true;
    R.mont = true;
    u256 z; u256_zero(z);
    R.x = to_mont(z);
    R.y = to_mont(z);
    R.z = to_mont(z);

    // Q = P as Jacobian (MONT)
    P256Point Q = toJac(Pstd);

    // classic double-and-add MSB->LSB (matches get_bit_be)
    for (int i=0;i<256;i++){
        uint32_t bit = get_bit_be(k,i);

        R = point_double(R);
        if(bit) R = point_add_jac(R,Q);
    }

    // ✅ Return affine STD
    return toAff(R);
}



P256Point scalar_mul_G(const uint8_t priv[32])
{
    if (scalar_is_zero(priv)) {
        P256Point R;
        R.inf = true;
        R.mont = false;
        return R;
    }
    return scalar_mul(priv, P256_G());
}

bool is_on_curve(const P256Point& P)
{
    if (P.inf) return true;

    // reject x>=p or y>=p
    if (u256_cmp(P.x, P256_P) >= 0) return false;
    if (u256_cmp(P.y, P256_P) >= 0) return false;

    u256 x = P.mont ? P.x : to_mont(P.x);
    u256 y = P.mont ? P.y : to_mont(P.y);
    u256 b = to_mont(P256_B);

    u256 y2 = fp_mul(y,y);
    u256 x2 = fp_mul(x,x);
    u256 x3 = fp_mul(x2,x);

    // rhs = x^3 - 3x + b
    u256 three_std = {{3,0,0,0,0,0,0,0}};
    u256 three = to_mont(three_std);

    u256 rhs = fp_sub(x3, fp_mul(three,x));
    rhs = fp_add(rhs,b);

    return u256_cmp(y2,rhs)==0;
}

// ------------------------------------------------------------
// TLS / NIST Interoperability
// ------------------------------------------------------------

/**
 * Decodes uncompressed TLS point (04||X||Y).
 * Internally stores coordinates in Montgomery domain.
 * Also validates curve equation in Montgomery domain.
 */
bool decode_tls_point(P256Point& out, const uint8_t* in, size_t len)
{
    if (len != 65 || in[0] != 0x04) return false;

    out.inf  = false;
    out.mont = false;
    out.x = bytes_be_to_u256(in + 1);
    out.y = bytes_be_to_u256(in + 33);
    return true;
}


/**
 * Encodes affine point (Montgomery) to TLS (std big endian).
 */
std::vector<uint8_t> encode_tls_point(const P256Point& P)
{
    std::vector<uint8_t> out(65);
    out[0]=0x04;
    u256_to_be(&out[1], P.x);
    u256_to_be(&out[33],P.y);
    return out;
}



/**
 * ECDH: returns Big-Endian X-coordinate of shared point.
 */
bool ecdh_shared_secret(uint8_t out32[32], const uint8_t priv[32], const P256Point& peer)
{
    if (!out32 || peer.inf) return false;

    // scalar_mul returns affine STD
    P256Point shared = scalar_mul(priv, peer);
    if (shared.inf) return false;

    u256_to_be(out32, shared.x);   // ✅ already STD
    return true;
}

} // namespace netplus
