#include <stdio.h>
#include <string.h>
#include <stdint.h>

/* Copy the exact P-256 constants and functions from ecc.c */
static const uint64_t P256_P[4] = {
    0xFFFFFFFFFFFFFFFFULL, 0x00000000FFFFFFFFULL,
    0x0000000000000000ULL, 0xFFFFFFFF00000001ULL
};

static const uint64_t P256_GX[4] = {
    0xF4A13945D898C296ULL, 0x77037D812DEB33A0ULL,
    0xF8BCE6E563A440F2ULL, 0x6B17D1F2E12C4247ULL
};

static const uint64_t P256_GY[4] = {
    0xCBB6406837BF51F5ULL, 0x2BCE33576B315ECEULL,
    0x8EE7EB4A7C0F9E16ULL, 0x4FE342E2FE1A7F9BULL
};

static const uint64_t THREE[4]={3,0,0,0};

static void bn_copy(uint64_t *a, const uint64_t *b)
{ a[0]=b[0]; a[1]=b[1]; a[2]=b[2]; a[3]=b[3]; }

static int bn_lt(const uint64_t *a, const uint64_t *b)
{
    for(int i=3;i>=0;i--) { if(a[i]<b[i]) return 1; if(a[i]>b[i]) return 0; }
    return 0;
}

static uint64_t bn_add(uint64_t *a, const uint64_t *b)
{
    uint64_t carry=0;
    for(int i=0;i<4;i++){
        uint64_t sum = a[i] + b[i] + carry;
        if(sum < a[i] || (carry && sum == a[i])) carry = 1;
        else carry = (sum < b[i]) ? 1 : 0;
        a[i] = sum;
    }
    return carry;
}

static uint64_t bn_sub(uint64_t *a, const uint64_t *b)
{
    uint64_t borrow=0;
    for(int i=0;i<4;i++){
        uint64_t diff = a[i] - b[i] - borrow;
        uint64_t new_borrow = (a[i] < b[i]) ? 1 : ((a[i] == b[i]) ? borrow : 0);
        borrow = new_borrow;
        a[i] = diff;
    }
    return borrow;
}

static void bn_mod_p256(uint64_t *a)
{
    while(bn_lt(P256_P,a)){bn_sub(a,P256_P);}
}

static void bn_add_mod(uint64_t *a, const uint64_t *b)
{
    uint64_t c=bn_add(a,b);
    if(c||bn_lt(P256_P,a)) bn_sub(a,P256_P);
}

static void bn_sub_mod(uint64_t *a, const uint64_t *b)
{
    uint64_t br=bn_sub(a,b);
    if(br) bn_add(a,P256_P);
}

static uint64_t mul64(uint64_t a, uint64_t b, uint64_t *lo)
{
    uint32_t a_lo = (uint32_t)a, a_hi = (uint32_t)(a >> 32);
    uint32_t b_lo = (uint32_t)b, b_hi = (uint32_t)(b >> 32);
    uint64_t ll = (uint64_t)a_lo * b_lo;
    uint64_t lh = (uint64_t)a_lo * b_hi;
    uint64_t hl = (uint64_t)a_hi * b_lo;
    uint64_t hh = (uint64_t)a_hi * b_hi;
    uint64_t mid = (ll >> 32) + (uint32_t)lh + (uint32_t)hl;
    *lo = (mid << 32) | (uint32_t)ll;
    return hh + (mid >> 32) + (lh >> 32) + (hl >> 32);
}

static void bn_mul_mod(uint64_t *c, const uint64_t *a, const uint64_t *b)
{
    uint64_t tmp[8]={0};
    for(int i=0;i<4;i++){
        uint64_t carry=0;
        for(int j=0;j<4;j++){
            uint64_t lo;
            uint64_t hi = mul64(a[i], b[j], &lo);
            uint64_t sum = tmp[i+j] + lo;
            uint64_t c1 = (sum < lo) ? 1 : 0;
            sum += carry;
            c1 += (sum < carry) ? 1 : 0;
            tmp[i+j] = sum;
            carry = hi + c1;
        }
        tmp[i+4]=carry;
    }

    uint32_t t[16];
    for(int i=0;i<4;i++){
        t[2*i]   = (uint32_t)(tmp[i]);
        t[2*i+1] = (uint32_t)(tmp[i]>>32);
    }
    for(int i=0;i<4;i++){
        t[8+2*i]   = (uint32_t)(tmp[4+i]);
        t[8+2*i+1] = (uint32_t)(tmp[4+i]>>32);
    }

    int64_t s[15];
    for(int i=0;i<15;i++) s[i]=0;
    for(int i=0;i<8;i++) s[i]=t[i];
    for(int k=8;k<16;k++){
        s[k-8]+=t[k]; s[k-5]-=t[k]; s[k-2]-=t[k]; s[k-1]+=t[k];
    }
    for(int k=8;k<15;k++){
        if(s[k]==0) continue;
        int64_t v=s[k]; s[k]=0;
        s[k-8]+=v; s[k-5]-=v; s[k-2]-=v; s[k-1]+=v;
    }
    for(int k=8;k<15;k++){
        if(s[k]==0) continue;
        int64_t v=s[k]; s[k]=0;
        s[k-8]+=v; s[k-5]-=v; s[k-2]-=v; s[k-1]+=v;
    }
    for(int i=0;i<14;i++){
        int64_t carry=s[i]>>32;
        s[i]&=0xFFFFFFFF;
        s[i+1]+=carry;
    }
    for(int k=8;k<15;k++){
        if(s[k]==0) continue;
        int64_t v=s[k]; s[k]=0;
        s[k-8]+=v; s[k-5]-=v; s[k-2]-=v; s[k-1]+=v;
    }
    for(int i=0;i<14;i++){
        int64_t carry=s[i]>>32;
        s[i]&=0xFFFFFFFF;
        s[i+1]+=carry;
    }

    for(int i=0;i<4;i++)
        c[i] = (uint64_t)s[2*i] | ((uint64_t)s[2*i+1] << 32);
    bn_mod_p256(c);
}

static void bn_sqr_mod(uint64_t *a)
{
    bn_mul_mod(a, a, a);
}

static void print_limbs(const char *label, uint64_t *v) {
    printf("%s: ", label);
    for(int i=3;i>=0;i--) printf("%016lx", v[i]);
    printf("\n");
}

int main() {
    uint64_t X2[4], X2_v2[4];
    
    /* Compute X^2 using mul_mod */
    bn_copy(X2, P256_GX);
    bn_mul_mod(X2, P256_GX, P256_GX);
    
    /* Compute X^2 using sqr_mod */
    bn_copy(X2_v2, P256_GX);
    bn_sqr_mod(X2_v2);
    
    print_limbs("X^2 via mul_mod", X2);
    print_limbs("X^2 via sqr_mod", X2_v2);
    
    printf("X^2 match: %s\n", memcmp(X2, X2_v2, 32) == 0 ? "YES" : "NO");
    
    /* Compare with expected */
    uint64_t expected_X2[4] = {0x002ae56c426b3f8cULL, 0x33b699495d694dd1ULL, 0x81819a5e0e3690d8ULL, 0x98f6b84d29bef2b2ULL};
    print_limbs("X^2 expected", expected_X2);
    printf("X^2 vs expected: %s\n", memcmp(X2, expected_X2, 32) == 0 ? "PASS" : "FAIL");
    
    /* Also compute 3*X^2 */
    uint64_t M1[4];
    bn_copy(M1, X2);
    bn_mul_mod(M1, M1, THREE);
    print_limbs("3*X^2", M1);
    
    /* Now compute Y^2 */
    uint64_t Y2[4];
    bn_copy(Y2, P256_GY);
    bn_sqr_mod(Y2);
    print_limbs("Y^2", Y2);
    
    uint64_t expected_Y2[4] = {0xbedd2cfc21c3ed91ULL, 0x498a9022a412b5e0ULL, 0x82149139979369feULL, 0x55df5d5850f47badULL};
    print_limbs("Y^2 expected", expected_Y2);
    printf("Y^2 vs expected: %s\n", memcmp(Y2, expected_Y2, 32) == 0 ? "PASS" : "FAIL");
    
    /* Diagnostic: test mul64 with known values */
    printf("\n--- mul64 diagnostics ---\n");
    {
        uint64_t lo, hi;
        /* Test 1: simple case */
        hi = mul64(5, 7, &lo);
        printf("5*7 = 0x%016lx%016lx (expected 0x00000000000000000000000000000023)\n", hi, lo);
        
        /* Test 2: all ones */
        uint64_t all = 0xFFFFFFFFFFFFFFFFULL;
        hi = mul64(all, all, &lo);
        printf("(2^64-1)^2 = 0x%016lx%016lx (expected 0xFFFFFFFFFFFFFFFE0000000000000001)\n", hi, lo);
        
        /* Test 3: large values */
        hi = mul64(0xF4A13945D898C296ULL, 0xF4A13945D898C296ULL, &lo);
        printf("Gx[0]*Gx[0] = 0x%016lx%016lx\n", hi, lo);
        
        /* Test 4: schoolbook product of Gx[0]*Gx[1] */
        hi = mul64(0xF4A13945D898C296ULL, 0x77037D812DEB33A0ULL, &lo);
        printf("Gx[0]*Gx[1] = 0x%016lx%016lx\n", hi, lo);
    }
    
    return 0;
}
