#include "crypto/ecc_u256.h"
#include <cstring>
#include <cstdio>

namespace netplus {

// ============================================================================
// P-256 prime:
// p = 0xffffffff00000001000000000000000000000000ffffffffffffffffffffffff
// internal u256 is LE words: w[0] = least significant 32 bits
// ============================================================================
const u256 P256_P = {{
    0xFFFFFFFFu, 0xFFFFFFFFu, 0xFFFFFFFFu, 0x00000000u,
    0x00000000u, 0x00000000u, 0x00000001u, 0xFFFFFFFFu
}};


static const u256 P256_R = {{
    0x00000001u,
    0x00000000u,
    0x00000000u,
    0xFFFFFFFFu,
    0xFFFFFFFFu,
    0xFFFFFFFFu,
    0xFFFFFFFEu,
    0x00000000u
}};

// R^2 mod p  (critical!)
//
// Correct:
// R^2 mod p = 0x00000004fffffffdfffffffffffffffefffffffbffffffff0000000000000003
//
const u256 P256_R2 = {{
    0x00000003,
    0x00000000,
    0xFFFFFFFF,
    0xFFFFFFFB,
    0xFFFFFFFE,
    0xFFFFFFFF,
    0xFFFFFFFD,
    0x00000004
}};

// n0 = -p^{-1} mod 2^32. Since p0 = 0xffffffff, inverse is 1, so n0=1.
static constexpr uint32_t P256_N0 = 1u;

// ----------------------------------------------------------------------------
// basic helpers
// ----------------------------------------------------------------------------
void u256_zero(u256& a) { std::memset(a.w, 0, sizeof(a.w)); }

bool u256_is_zero(const u256& a) {
    for (int i = 0; i < 8; i++) if (a.w[i]) return false;
    return true;
}

int u256_cmp(const u256& a, const u256& b) {
    for (int i = 7; i >= 0; --i) {
        if (a.w[i] < b.w[i]) return -1;
        if (a.w[i] > b.w[i]) return 1;
    }
    return 0;
}

u256 u256_add_raw(const u256& a, const u256& b, uint32_t& carry) {
    u256 r;
    uint64_t c = 0;
    for (int i = 0; i < 8; i++) {
        uint64_t s = (uint64_t)a.w[i] + b.w[i] + c;
        r.w[i] = (uint32_t)s;
        c = s >> 32;
    }
    carry = (uint32_t)c;
    return r;
}

u256 u256_sub_raw(const u256& a, const u256& b, uint32_t& borrow) {
    u256 r;
    uint64_t br = 0;
    for (int i = 0; i < 8; i++) {
        uint64_t ai = (uint64_t)a.w[i];
        uint64_t bi = (uint64_t)b.w[i];
        uint64_t d  = ai - bi - br;
        r.w[i] = (uint32_t)d;
        br = (d >> 63) & 1;
    }
    borrow = (uint32_t)br;
    return r;
}

// ----------------------------------------------------------------------------
// BE/LE conversion
// ----------------------------------------------------------------------------
u256 u256_from_be(const uint8_t be32[32]) {
    u256 r;
    for (int i = 0; i < 8; i++) {
        int o = 28 - 4*i;
        r.w[i] =
            ((uint32_t)be32[o+0] << 24) |
            ((uint32_t)be32[o+1] << 16) |
            ((uint32_t)be32[o+2] << 8)  |
            ((uint32_t)be32[o+3] << 0);
    }
    return r;
}

void u256_to_be(uint8_t out32[32], const u256& a) {
    for (int i = 0; i < 8; i++) {
        uint32_t w = a.w[i];
        int o = 28 - 4*i;
        out32[o+0] = (uint8_t)(w >> 24);
        out32[o+1] = (uint8_t)(w >> 16);
        out32[o+2] = (uint8_t)(w >> 8);
        out32[o+3] = (uint8_t)(w >> 0);
    }
}

u256 bytes_be_to_u256(const uint8_t in[32]) {
    u256 r;
    for (int i = 0; i < 8; i++) {
        int o = 28 - 4*i;   // ✅ same mapping as u256_to_be
        r.w[i] =
            ((uint32_t)in[o+0] << 24) |
            ((uint32_t)in[o+1] << 16) |
            ((uint32_t)in[o+2] <<  8) |
            ((uint32_t)in[o+3] <<  0);
    }
    return r;
}


// ============================================================================
// 256x256 -> 512 multiply (schoolbook, correct carry), NO __int128
// ============================================================================
u512 u256_mul_raw(const u256& a, const u256& b) {
    u512 out;
    for (int i=0; i<16; i++) out.w[i] = 0;

    for (int i=0; i<8; i++){
        uint64_t c = 0;
        // Innere Schleife: Multiplikation und Addition
        for (int j=0; j<8; j++){
            uint64_t uv = (uint64_t)out.w[i+j] + (uint64_t)a.w[i]*b.w[j] + c;
            out.w[i+j] = (uint32_t)uv;
            c = uv >> 32;
        }

        // Übertrag auf die nächsten Wörter propagieren
        // Wir nutzen hier uint64_t, um Überläufe sicher zu handhaben
        int k = i + 8;
        while (c > 0 && k < 16) {
            uint64_t uv = (uint64_t)out.w[k] + c;
            out.w[k] = (uint32_t)uv;
            c = uv >> 32; // Wird nur > 0 sein, wenn out.w[k] überlief
            k++;
        }
    }
    return out;
}

// input: 64 bytes Big Endian (Network order)
// output: u512 with w[0] = least significant 32-bit word
static inline u512 u512_from_be64(const uint8_t be64[64]) {
    u512 r;
    for (int i = 0; i < 16; i++) {
        // be64[0..3] = MSW, be64[60..63] = LSW
        int o = 4 * i; // read MSW -> LSW in increasing offset
        uint32_t w =
            ((uint32_t)be64[o + 0] << 24) |
            ((uint32_t)be64[o + 1] << 16) |
            ((uint32_t)be64[o + 2] <<  8) |
            ((uint32_t)be64[o + 3] <<  0);

        // store into LE word array:
        r.w[15 - i] = w;
    }
    return r;
}


// ============================================================================
// Montgomery REDC for P-256 (8 limbs, 32-bit, no __int128)
// This version is stable and matches classic HAC REDC.
u256 mont_mul(const u256& a, const u256& b) {
    uint32_t t[16] = {0};
    uint32_t extra = 0; // overflow beyond t[15]

    // ---------------------------------------------------------
    // Step 1: t = a*b  (full 512-bit, propagate carry fully)
    // ---------------------------------------------------------
    for (int i=0; i<8; i++) {
        uint64_t carry = 0;
        for (int j=0; j<8; j++) {
            uint64_t uv = (uint64_t)t[i+j]
                        + (uint64_t)a.w[i] * b.w[j]
                        + carry;
            t[i+j] = (uint32_t)uv;
            carry  = uv >> 32;
        }

        int k = i + 8;
        while (carry) {
            if (k < 16) {
                uint64_t uv = (uint64_t)t[k] + carry;
                t[k] = (uint32_t)uv;
                carry = uv >> 32;
                k++;
            } else {
                extra += (uint32_t)carry; // should be only 1
                break;
            }
        }
    }

    // ---------------------------------------------------------
    // Step 2: Montgomery REDC
    // ---------------------------------------------------------
    for (int i=0; i<8; i++) {
        uint32_t m = t[i] * P256_N0; // N0=1 for P-256

        uint64_t carry = 0;
        for (int j=0; j<8; j++) {
            uint64_t uv = (uint64_t)t[i+j]
                        + (uint64_t)m * P256_P.w[j]
                        + carry;
            t[i+j] = (uint32_t)uv;
            carry  = uv >> 32;
        }

        int k = i + 8;
        while (carry) {
            if (k < 16) {
                uint64_t uv = (uint64_t)t[k] + carry;
                t[k] = (uint32_t)uv;
                carry = uv >> 32;
                k++;
            } else {
                extra += (uint32_t)carry;
                break;
            }
        }
    }

    // ---------------------------------------------------------
    // Step 3: r = t[8..15], normalize
    // ---------------------------------------------------------
    u256 r;
    for (int i=0; i<8; i++) r.w[i] = t[i+8];

    if (extra || u256_cmp(r, P256_P) >= 0) {
        uint32_t br=0;
        r = u256_sub_raw(r, P256_P, br);
    }
    if (u256_cmp(r, P256_P) >= 0) {
        uint32_t br=0;
        r = u256_sub_raw(r, P256_P, br);
    }

    return r;
}

// ----------------------------------------------------------------------------
// to/from Montgomery
// ----------------------------------------------------------------------------
u256 to_mont(const u256& a) {
    return mont_mul(a, P256_R2);
}

u256 from_mont(const u256& aR) {
    u256 one; u256_zero(one); one.w[0]=1;
    return mont_mul(aR, one); // NUR wenn mont_mul als REDC(t) implementiert ist
}

// ----------------------------------------------------------------------------
// field ops (Montgomery domain)
// ----------------------------------------------------------------------------
u256 fp_add(const u256& a, const u256& b) {
    uint32_t carry = 0;
    u256 r = u256_add_raw(a, b, carry);
    if (carry || u256_cmp(r, P256_P) >= 0) {
        uint32_t br = 0;
        r = u256_sub_raw(r, P256_P, br);
    }
    return r;
}

u256 fp_sub(const u256& a, const u256& b) {
    uint32_t br = 0;
    u256 r = u256_sub_raw(a, b, br);
    if (br) {
        uint32_t c = 0;
        r = u256_add_raw(r, P256_P, c);
    }
    return r;
}

u256 fp_mul(const u256& a, const u256& b) {
    return mont_mul(a, b);
}

u256 fp_sqr(const u256& a) {
    return mont_mul(a, a);
}

// ----------------------------------------------------------------------------
// bit access
// ----------------------------------------------------------------------------
uint32_t get_bit_be(const uint8_t priv[32], int bit_index) {
    // i=0 -> MSB of priv[0]
    int byte = bit_index >> 3;
    int bit  = 7 - (bit_index & 7);
    return (priv[byte] >> bit) & 1;
}

u512 u512_from_be(const uint8_t in[64]) {
    u512 r;
    for (int i=0;i<16;i++){
        int o = 60 - 4*i;
        r.w[i] =
            ((uint32_t)in[o+0] << 24) |
            ((uint32_t)in[o+1] << 16) |
            ((uint32_t)in[o+2] <<  8) |
            ((uint32_t)in[o+3] <<  0);
    }
    return r; // r.w[0] is LSW
}

u256 mod_p256(const u512& in)
{
    const uint32_t* a = in.w; // little-endian 32-bit limbs

    int64_t t[16] = {0};

    // t = full 512-bit input (L + H*2^256)
    for(int i=0;i<16;i++) t[i] = (int64_t)a[i];

    auto add64 = [&](int idx, int64_t v){
        if ((unsigned)idx < 16) t[idx] += v;
    };
    auto sub64 = [&](int idx, int64_t v){
        if ((unsigned)idx < 16) t[idx] -= v;
    };

    // signed normalize across all 16 limbs
    auto normalize16 = [&](){
        int64_t carry = 0;
        for(int i=0;i<16;i++){
            t[i] += carry;
            carry = t[i] >> 32;
            t[i] &= 0xffffffffLL;
        }
        return carry; // should be small
    };

    // fold one high-limb word k at shift sh (word-shift)
    // k * 2^(256 + 32*sh)  with 2^256 ≡ 1 + 2^224 - 2^192 - 2^96
    auto fold_k_at_shift = [&](int64_t k, int sh){
        add64(sh + 0, k);  // +k
        add64(sh + 7, k);  // +k<<224
        sub64(sh + 6, k);  // -k<<192
        sub64(sh + 3, k);  // -k<<96
    };

    auto fold_hi = [&](int i){
        int64_t k = t[i];
        if(k == 0) return;
        t[i] = 0;
        int sh = i - 8;
        fold_k_at_shift(k, sh);
    };

    // First fold pass
    for(int i=15;i>=8;i--) fold_hi(i);
    int64_t top = normalize16();

    // Second fold pass (needed because fold can create new highs)
    for(int i=15;i>=8;i--) fold_hi(i);
    top += normalize16();

    // Fold remaining carry (treat as limb i=16 => sh=8)
    // top*2^512 = top * (2^256)^2 → folding top at sh=8 is enough once, then re-fold highs.
    while(top != 0){
        int64_t k = top;
        top = 0;
        fold_k_at_shift(k, 8);
        top += normalize16();
        for(int i=15;i>=8;i--) fold_hi(i);
        top += normalize16();
    }

    // extract low 256 bits
    u256 r;
    for(int i=0;i<8;i++) r.w[i] = (uint32_t)t[i];

    // constant-time subtract p twice (safe)
    auto ct_sub_p = [&](){
        uint32_t br=0;
        u256 diff = u256_sub_raw(r, P256_P, br);
        uint32_t mask = (uint32_t)0 - (uint32_t)(br ^ 1);
        for(int i=0;i<8;i++){
            r.w[i] = (r.w[i] & ~mask) | (diff.w[i] & mask);
        }
    };

    ct_sub_p();
    ct_sub_p();

    return r;
}

u256 mod_p256(const uint8_t be64[64]) {
    u512 x = u512_from_be64(be64);
    return mod_p256(x);
}

// ----------------------------------------------------------------------------
// exponentiation/inversion
// ----------------------------------------------------------------------------
static u256 fp_pow(const u256& a_mont, const u256& e_normal) {
    u256 one; u256_zero(one); one.w[0] = 1;
    u256 r = to_mont(one);

    for (int i = 7; i >= 0; i--) {
        for (int bit = 31; bit >= 0; bit--) {
            r = fp_sqr(r);
            if ((e_normal.w[i] >> bit) & 1) {
                r = fp_mul(r, a_mont);
            }
        }
    }
    return r;
}

u256 fp_inv(const u256& a_mont) {
    u256 exp = P256_P;
    u256 two; u256_zero(two); two.w[0] = 2;
    uint32_t br = 0;
    exp = u256_sub_raw(exp, two, br);
    return fp_pow(a_mont, exp);
}

// ============================================================================
// P-256 group order n for scalar field arithmetic (ECDSA)
// n = 0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551
// ============================================================================
static const u256 SCALAR_N = {{
    0xFC632551, // w[0] LSW
    0xF3B9CAC2,
    0xA7179E84,
    0xBCE6FAAD,
    0xFFFFFFFF,
    0xFFFFFFFF,
    0x00000000,
    0xFFFFFFFF  // w[7] MSW
}};

// Reduction mod n (group order) - used for ECDSA scalar arithmetic
// Uses schoolbook division since n doesn't have a special form like p
u256 mod_n256(const u512& in) {
    // For a 512-bit value, we do trial subtraction of n*2^i for i from 256 down to 0
    // This is slow but correct. For production, use Barrett reduction.
    
    // Work with extended precision
    uint64_t t[16];
    for (int i = 0; i < 16; i++) t[i] = in.w[i];
    
    // Reduce: while value >= n * 2^256, subtract n * 2^256
    // Then while value >= n * 2^224, subtract n * 2^224, etc.
    
    // Simple approach: convert to u256, handle overflow with subtraction
    // First, check if high 256 bits are non-zero
    bool hasHigh = false;
    for (int i = 8; i < 16; i++) if (t[i] != 0) { hasHigh = true; break; }
    
    if (!hasHigh) {
        // Just low 256 bits - simple reduction
        u256 r;
        for (int i = 0; i < 8; i++) r.w[i] = (uint32_t)t[i];
        while (u256_cmp(r, SCALAR_N) >= 0) {
            uint32_t br = 0;
            r = u256_sub_raw(r, SCALAR_N, br);
        }
        return r;
    }
    
    // Full 512-bit reduction using repeated subtraction
    // This is O(n) in the worst case but works for random inputs
    
    // Compute q ≈ floor(in / n) using high bits estimation
    // Then r = in - q*n and adjust
    
    // For simplicity, use a loop that processes word by word from high to low
    // Each iteration reduces the value by subtracting n shifted appropriately
    
    for (int shift = 256; shift >= 0; shift -= 32) {
        // Check if we can subtract n << shift
        int wordShift = shift / 32;
        
        while (true) {
            // Check if t >= n << shift
            bool canSubtract = false;
            
            // Compare t with n << wordShift
            int topWord = wordShift + 7;
            if (topWord < 16) {
                // Check from top
                bool greater = false;
                bool equal = true;
                for (int i = 15; i >= 0; i--) {
                    uint64_t nShifted = 0;
                    int ni = i - wordShift;
                    if (ni >= 0 && ni < 8) nShifted = SCALAR_N.w[ni];
                    
                    if (t[i] > nShifted) { greater = true; equal = false; break; }
                    if (t[i] < nShifted) { greater = false; equal = false; break; }
                }
                canSubtract = greater || equal;
            }
            
            if (!canSubtract) break;
            
            // Subtract n << wordShift from t
            int64_t borrow = 0;
            for (int i = 0; i < 16; i++) {
                uint64_t nShifted = 0;
                int ni = i - wordShift;
                if (ni >= 0 && ni < 8) nShifted = SCALAR_N.w[ni];
                
                int64_t diff = (int64_t)t[i] - (int64_t)nShifted - borrow;
                if (diff < 0) {
                    diff += 0x100000000LL;
                    borrow = 1;
                } else {
                    borrow = 0;
                }
                t[i] = (uint64_t)diff;
            }
        }
    }
    
    // Result is now in t[0..7]
    u256 r;
    for (int i = 0; i < 8; i++) r.w[i] = (uint32_t)t[i];
    
    // Final normalization
    while (u256_cmp(r, SCALAR_N) >= 0) {
        uint32_t br = 0;
        r = u256_sub_raw(r, SCALAR_N, br);
    }
    
    return r;
}

// Scalar field multiplication: (a * b) mod n
u256 scalar_mul_mod_n(const u256& a, const u256& b) {
    u512 prod = u256_mul_raw(a, b);
    return mod_n256(prod);
}

// Scalar field addition: (a + b) mod n
u256 scalar_add_mod_n(const u256& a, const u256& b) {
    uint32_t carry = 0;
    u256 r = u256_add_raw(a, b, carry);
    
    // If carry or r >= n, subtract n
    if (carry || u256_cmp(r, SCALAR_N) >= 0) {
        uint32_t br = 0;
        r = u256_sub_raw(r, SCALAR_N, br);
    }
    return r;
}

// Scalar field inversion using Fermat's little theorem: a^-1 = a^(n-2) mod n
u256 scalar_inv_mod_n(const u256& a) {
    // n-2 for SCALAR_N
    static const u256 N_MINUS_2 = {{
        0xFC63254F, // w[0] = 0xFC632551 - 2
        0xF3B9CAC2,
        0xA7179E84,
        0xBCE6FAAD,
        0xFFFFFFFF,
        0xFFFFFFFF,
        0x00000000,
        0xFFFFFFFF
    }};
    
    // Square-and-multiply: result = a^(n-2) mod n
    u256 result;
    u256_zero(result);
    result.w[0] = 1;
    
    u256 base = a;
    
    for (int i = 0; i < 8; i++) {
        for (int bit = 0; bit < 32; bit++) {
            if ((N_MINUS_2.w[i] >> bit) & 1) {
                result = scalar_mul_mod_n(result, base);
            }
            base = scalar_mul_mod_n(base, base);
        }
    }
    
    return result;
}

} // namespace netplus
