const std = @import("std");
const Secp256k1 = std.crypto.ecc.Secp256k1;
const StdFe = Secp256k1.Fe;

// 5×52-bit unsaturated limb representation (libsecp256k1 style).
// Each limb holds up to 52 bits of value with 12 bits of headroom in u64.
// p = 2^256 - 2^32 - 977, so R = 2^256 mod p = 0x1000003D1.
const M: u64 = 0xFFFFFFFFFFFFF; // 52-bit mask
const R: u64 = 0x1000003D10; // 2^256 mod p (shifted by 4 for limb4's 48-bit width)

// secp256k1 prime p in 5×52
const P0: u64 = 0xFFFFEFFFFFC2F;
const P4: u64 = 0xFFFFFFFFFFFF; // 48 bits (limb 4 is only 48 bits wide)

/// Optimized secp256k1 field element: 5 × 52-bit limbs in 64-bit words.
///
/// 12 bits of headroom per limb (4 bits in MSW) enables carry-free add/sub
/// for magnitudes up to ~4096. Mul/sq use libsecp256k1's interleaved
/// reduction with R = 2^256 mod p, producing fully reduced output.
pub const Fe = struct {
    n: [5]u64,

    pub const zero: Fe = .{ .n = .{ 0, 0, 0, 0, 0 } };
    pub const one: Fe = .{ .n = .{ 1, 0, 0, 0, 0 } };

    /// Pack a 32-byte big-endian value into 5×52-bit limbs.
    pub fn fromBytes(b: [32]u8, comptime endian: std.builtin.Endian) Fe {
        const s = if (endian == .big) b else blk: {
            var rev: [32]u8 = undefined;
            for (0..32) |i| rev[i] = b[31 - i];
            break :blk rev;
        };
        // s is big-endian: s[0] is MSB, s[31] is LSB
        // limb 0 = bits 0..51 (least significant)
        // limb 4 = bits 208..255 (most significant, 48 bits)
        return .{ .n = .{
            @as(u64, s[31]) | @as(u64, s[30]) << 8 | @as(u64, s[29]) << 16 |
                @as(u64, s[28]) << 24 | @as(u64, s[27]) << 32 | @as(u64, s[26]) << 40 |
                (@as(u64, s[25]) & 0xF) << 48,
            @as(u64, s[25]) >> 4 | @as(u64, s[24]) << 4 | @as(u64, s[23]) << 12 |
                @as(u64, s[22]) << 20 | @as(u64, s[21]) << 28 | @as(u64, s[20]) << 36 |
                @as(u64, s[19]) << 44,
            @as(u64, s[18]) | @as(u64, s[17]) << 8 | @as(u64, s[16]) << 16 |
                @as(u64, s[15]) << 24 | @as(u64, s[14]) << 32 | @as(u64, s[13]) << 40 |
                (@as(u64, s[12]) & 0xF) << 48,
            @as(u64, s[12]) >> 4 | @as(u64, s[11]) << 4 | @as(u64, s[10]) << 12 |
                @as(u64, s[9]) << 20 | @as(u64, s[8]) << 28 | @as(u64, s[7]) << 36 |
                @as(u64, s[6]) << 44,
            @as(u64, s[5]) | @as(u64, s[4]) << 8 | @as(u64, s[3]) << 16 |
                @as(u64, s[2]) << 24 | @as(u64, s[1]) << 32 | @as(u64, s[0]) << 40,
        } };
    }

    /// Unpack 5×52-bit limbs to a 32-byte big-endian value.
    /// The field value MUST be normalized first.
    pub fn toBytes(self: Fe, comptime endian: std.builtin.Endian) [32]u8 {
        const n = self.n;
        var b: [32]u8 = undefined;
        // limb 0: bits 0..51
        b[31] = @truncate(n[0]);
        b[30] = @truncate(n[0] >> 8);
        b[29] = @truncate(n[0] >> 16);
        b[28] = @truncate(n[0] >> 24);
        b[27] = @truncate(n[0] >> 32);
        b[26] = @truncate(n[0] >> 40);
        // limb 1: bits 52..103
        b[25] = @truncate((n[0] >> 48) | (n[1] << 4));
        b[24] = @truncate(n[1] >> 4);
        b[23] = @truncate(n[1] >> 12);
        b[22] = @truncate(n[1] >> 20);
        b[21] = @truncate(n[1] >> 28);
        b[20] = @truncate(n[1] >> 36);
        b[19] = @truncate(n[1] >> 44);
        // limb 2: bits 104..155
        b[18] = @truncate(n[2]);
        b[17] = @truncate(n[2] >> 8);
        b[16] = @truncate(n[2] >> 16);
        b[15] = @truncate(n[2] >> 24);
        b[14] = @truncate(n[2] >> 32);
        b[13] = @truncate(n[2] >> 40);
        // limb 3: bits 156..207
        b[12] = @truncate((n[2] >> 48) | (n[3] << 4));
        b[11] = @truncate(n[3] >> 4);
        b[10] = @truncate(n[3] >> 12);
        b[9] = @truncate(n[3] >> 20);
        b[8] = @truncate(n[3] >> 28);
        b[7] = @truncate(n[3] >> 36);
        b[6] = @truncate(n[3] >> 44);
        // limb 4: bits 208..255 (48 bits)
        b[5] = @truncate(n[4]);
        b[4] = @truncate(n[4] >> 8);
        b[3] = @truncate(n[4] >> 16);
        b[2] = @truncate(n[4] >> 24);
        b[1] = @truncate(n[4] >> 32);
        b[0] = @truncate(n[4] >> 40);
        if (endian == .little) {
            var rev: [32]u8 = undefined;
            for (0..32) |i| rev[i] = b[31 - i];
            return rev;
        }
        return b;
    }

    /// Carry propagation + mod-p reduction.
    pub fn normalize(self: Fe) Fe {
        var t0 = self.n[0];
        var t1 = self.n[1];
        var t2 = self.n[2];
        var t3 = self.n[3];
        var t4 = self.n[4];

        // propagate carries
        var m = t0 >> 52;
        t0 &= M;
        t1 += m;
        m = t1 >> 52;
        t1 &= M;
        t2 += m;
        m = t2 >> 52;
        t2 &= M;
        t3 += m;
        m = t3 >> 52;
        t3 &= M;
        t4 += m;

        // reduce: if t4 overflows 48 bits, fold back via p
        m = t4 >> 48;
        t4 &= 0xFFFFFFFFFFFF;
        t0 += m * 0x1000003D1;
        // re-propagate
        m = t0 >> 52;
        t0 &= M;
        t1 += m;
        m = t1 >> 52;
        t1 &= M;
        t2 += m;
        m = t2 >> 52;
        t2 &= M;
        t3 += m;
        m = t3 >> 52;
        t3 &= M;
        t4 += m;

        // conditional subtract p
        // check if value >= p
        const ge: u64 = if (t4 == 0xFFFFFFFFFFFF and
            t3 == M and t2 == M and t1 == M and
            t0 >= P0) 1 else 0;
        t0 += ge * 0x1000003D1;
        m = t0 >> 52;
        t0 &= M;
        t1 += m;
        m = t1 >> 52;
        t1 &= M;
        t2 += m;
        m = t2 >> 52;
        t2 &= M;
        t3 += m;
        m = t3 >> 52;
        t3 &= M;
        t4 += m;
        t4 &= 0xFFFFFFFFFFFF;

        return .{ .n = .{ t0, t1, t2, t3, t4 } };
    }

    /// Returns true if the field value equals zero (after normalizing).
    pub fn isZero(self: Fe) bool {
        const f = self.normalize();
        return (f.n[0] | f.n[1] | f.n[2] | f.n[3] | f.n[4]) == 0;
    }

    /// Returns true if two field values are equal (after normalizing both).
    pub fn equivalent(self: Fe, other: Fe) bool {
        const a = self.normalize();
        const b = other.normalize();
        return ((a.n[0] ^ b.n[0]) | (a.n[1] ^ b.n[1]) | (a.n[2] ^ b.n[2]) |
            (a.n[3] ^ b.n[3]) | (a.n[4] ^ b.n[4])) == 0;
    }

    /// Carry-free addition.
    pub fn add(self: Fe, other: Fe) Fe {
        return .{ .n = .{
            self.n[0] + other.n[0],
            self.n[1] + other.n[1],
            self.n[2] + other.n[2],
            self.n[3] + other.n[3],
            self.n[4] + other.n[4],
        } };
    }

    /// Negate: result = (magnitude+1)*p - self.
    pub fn neg(self: Fe, magnitude: u32) Fe {
        const m: u64 = @as(u64, magnitude) + 1;
        return .{ .n = .{
            m * P0 - self.n[0],
            m * M - self.n[1],
            m * M - self.n[2],
            m * M - self.n[3],
            m * P4 - self.n[4],
        } };
    }

    /// sub = self + (-other).
    pub fn subMag(self: Fe, other: Fe, other_mag: u32) Fe {
        return self.add(other.neg(other_mag));
    }

    /// sub with auto-normalize.
    pub fn sub(self: Fe, other: Fe) Fe {
        return self.add(other.normalize().neg(1));
    }

    /// Double via carry-free addition.
    pub fn dbl(self: Fe) Fe {
        return self.add(self);
    }

    /// Multiply by small integer.
    pub fn mulInt(self: Fe, val: u8) Fe {
        const v: u64 = val;
        return .{ .n = .{
            self.n[0] * v, self.n[1] * v, self.n[2] * v,
            self.n[3] * v, self.n[4] * v,
        } };
    }

    /// 5×52 multiplication with interleaved secp256k1 reduction.
    /// Ported from libsecp256k1's secp256k1_fe_mul_inner.
    /// Uses u128 accumulators with R = 0x1000003D10.
    pub fn mul(self: Fe, other: Fe) Fe {
        const a = self.n;
        const b = other.n;

        // column 3: d = p3
        var d: u128 = @as(u128, a[0]) * b[3] + @as(u128, a[1]) * b[2] +
            @as(u128, a[2]) * b[1] + @as(u128, a[3]) * b[0];

        // column 8: c = p8
        var c: u128 = @as(u128, a[4]) * b[4];

        // reduce p8 into d: d += low64(c) * R, c >>= 64
        d += @as(u128, R) * @as(u64, @truncate(c));
        c >>= 64;

        const t3 = @as(u64, @truncate(d)) & M;
        d >>= 52;

        // column 4: d += p4
        d += @as(u128, a[0]) * b[4] + @as(u128, a[1]) * b[3] +
            @as(u128, a[2]) * b[2] + @as(u128, a[3]) * b[1] + @as(u128, a[4]) * b[0];
        // reduce remaining c
        d += @as(u128, R << 12) * @as(u64, @truncate(c));

        var t4 = @as(u64, @truncate(d)) & M;
        d >>= 52;
        const tx = t4 >> 48;
        t4 &= (M >> 4);

        // column 0: c = p0
        c = @as(u128, a[0]) * b[0];

        // column 5: d += p5
        d += @as(u128, a[1]) * b[4] + @as(u128, a[2]) * b[3] +
            @as(u128, a[3]) * b[2] + @as(u128, a[4]) * b[1];
        var u_0 = @as(u64, @truncate(d)) & M;
        d >>= 52;

        u_0 = (u_0 << 4) | tx;
        c += @as(u128, u_0) * (R >> 4);

        const r0 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        // column 1: c += p1
        c += @as(u128, a[0]) * b[1] + @as(u128, a[1]) * b[0];

        // column 6: d += p6
        d += @as(u128, a[2]) * b[4] + @as(u128, a[3]) * b[3] + @as(u128, a[4]) * b[2];
        // reduce low 52 bits of d into c
        c += @as(u128, R) * (@as(u64, @truncate(d)) & M);
        d >>= 52;

        const r1 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        // column 2: c += p2
        c += @as(u128, a[0]) * b[2] + @as(u128, a[1]) * b[1] + @as(u128, a[2]) * b[0];

        // column 7: d += p7
        d += @as(u128, a[3]) * b[4] + @as(u128, a[4]) * b[3];
        // reduce: c += low64(d) * R, d >>= 64
        c += @as(u128, R) * @as(u64, @truncate(d));
        d >>= 64;

        const r2 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        // finalize r3, r4
        c += @as(u128, R << 12) * @as(u64, @truncate(d));
        c += t3;

        const r3 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        const r4 = @as(u64, @truncate(c)) + t4;

        return .{ .n = .{ r0, r1, r2, r3, r4 } };
    }

    /// 5×52 squaring with interleaved secp256k1 reduction.
    /// Ported from libsecp256k1's secp256k1_fe_sqr_inner.
    pub fn sq(self: Fe) Fe {
        const a = self.n;

        // column 3: d = p3 = 2*a0*a3 + 2*a1*a2
        var d: u128 = @as(u128, a[0] * 2) * a[3] + @as(u128, a[1] * 2) * a[2];

        // column 8: c = p8 = a4*a4
        var c: u128 = @as(u128, a[4]) * a[4];

        d += @as(u128, R) * @as(u64, @truncate(c));
        c >>= 64;

        const t3 = @as(u64, @truncate(d)) & M;
        d >>= 52;

        // column 4: d += p4 = a0*2*a4 + 2*a1*a3 + a2*a2
        const a4x2 = a[4] * 2;
        d += @as(u128, a[0]) * a4x2 + @as(u128, a[1] * 2) * a[3] + @as(u128, a[2]) * a[2];
        d += @as(u128, R << 12) * @as(u64, @truncate(c));

        var t4 = @as(u64, @truncate(d)) & M;
        d >>= 52;
        const tx = t4 >> 48;
        t4 &= (M >> 4);

        // column 0: c = p0 = a0*a0
        c = @as(u128, a[0]) * a[0];

        // column 5: d += p5 = a1*2*a4 + 2*a2*a3
        d += @as(u128, a[1]) * a4x2 + @as(u128, a[2] * 2) * a[3];
        var u_0 = @as(u64, @truncate(d)) & M;
        d >>= 52;

        u_0 = (u_0 << 4) | tx;
        c += @as(u128, u_0) * (R >> 4);

        const r0 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        // column 1: c += p1 = 2*a0*a1
        const a0x2 = a[0] * 2;
        c += @as(u128, a0x2) * a[1];

        // column 6: d += p6 = a2*2*a4 + a3*a3
        d += @as(u128, a[2]) * a4x2 + @as(u128, a[3]) * a[3];
        c += @as(u128, R) * (@as(u64, @truncate(d)) & M);
        d >>= 52;

        const r1 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        // column 2: c += p2 = 2*a0*a2 + a1*a1
        c += @as(u128, a0x2) * a[2] + @as(u128, a[1]) * a[1];

        // column 7: d += p7 = a3*2*a4
        d += @as(u128, a[3]) * a4x2;
        c += @as(u128, R) * @as(u64, @truncate(d));
        d >>= 64;

        const r2 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        c += @as(u128, R << 12) * @as(u64, @truncate(d));
        c += t3;

        const r3 = @as(u64, @truncate(c)) & M;
        c >>= 52;

        const r4 = @as(u64, @truncate(c)) + t4;

        return .{ .n = .{ r0, r1, r2, r3, r4 } };
    }

    /// Comptime conversion from integer literal to Fe.
    pub fn fromInt(comptime val: u256) Fe {
        var b: [32]u8 = undefined;
        std.mem.writeInt(u256, &b, val, .big);
        return fromBytes(b, .big);
    }

    /// Modular inverse via Fermat's little theorem: a^(p-2) mod p.
    pub fn invert(self: Fe) Fe {
        const a = self;
        const a2 = a.sq().mul(a);
        const a3 = a2.sq().mul(a);
        const a6 = sqn(a3, 3).mul(a3);
        const a9 = sqn(a6, 3).mul(a3);
        const a11 = sqn(a9, 2).mul(a2);
        const a22 = sqn(a11, 11).mul(a11);
        const a44 = sqn(a22, 22).mul(a22);
        const a88 = sqn(a44, 44).mul(a44);
        const a176 = sqn(a88, 88).mul(a88);
        const a220 = sqn(a176, 44).mul(a44);
        const a223 = sqn(a220, 3).mul(a3);

        var r = sqn(a223, 23).mul(a22);
        r = sqn(r, 5).mul(a);
        r = sqn(r, 3).mul(a2);
        r = sqn(r, 2).mul(a);
        return r;
    }

    fn sqn(f: Fe, comptime n: usize) Fe {
        var r = f;
        for (0..n) |_| r = r.sq();
        return r;
    }

    /// Convert from stdlib Fe to Fe via byte round-trip.
    pub fn fromStdlib(fe: StdFe) Fe {
        return fromBytes(fe.toBytes(.big), .big);
    }

    /// Convert to stdlib Fe via byte round-trip.
    pub fn toStdlib(self: Fe) StdFe {
        return StdFe.fromBytes(self.normalize().toBytes(.big), .big) catch unreachable;
    }
};

// ============================================================
// tests
// ============================================================

const testing = std.testing;

fn stdFeMul(a: StdFe, b: StdFe) StdFe {
    return a.mul(b);
}

fn stdFeSq(a: StdFe) StdFe {
    return a.sq();
}

test "fromBytes/toBytes round-trip" {
    const gx_bytes = Secp256k1.basePoint.affineCoordinates().x.toBytes(.big);
    const fe = Fe.fromBytes(gx_bytes, .big);
    const out = fe.normalize().toBytes(.big);
    try testing.expectEqualSlices(u8, &gx_bytes, &out);
}

test "fromBytes/toBytes little-endian round-trip" {
    const gx_bytes = Secp256k1.basePoint.affineCoordinates().x.toBytes(.little);
    const fe = Fe.fromBytes(gx_bytes, .little);
    const out = fe.normalize().toBytes(.little);
    try testing.expectEqualSlices(u8, &gx_bytes, &out);
}

test "zero and one" {
    try testing.expect(Fe.zero.isZero());
    try testing.expect(!Fe.one.isZero());
    try testing.expect(Fe.one.equivalent(Fe.one));
    try testing.expect(!Fe.zero.equivalent(Fe.one));
}

test "add matches stdlib" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const b = Fe.fromStdlib(g.y);
    const result = a.add(b).normalize().toBytes(.big);
    const expected = g.x.add(g.y).toBytes(.big);
    try testing.expectEqualSlices(u8, &expected, &result);
}

test "sub matches stdlib" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const b = Fe.fromStdlib(g.y);
    const result = a.sub(b).normalize().toBytes(.big);
    const expected = g.x.sub(g.y).toBytes(.big);
    try testing.expectEqualSlices(u8, &expected, &result);
}

test "neg matches stdlib" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const result = a.neg(1).normalize().toBytes(.big);
    const expected = StdFe.zero.sub(g.x).toBytes(.big);
    try testing.expectEqualSlices(u8, &expected, &result);
}

test "mul matches stdlib" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const b = Fe.fromStdlib(g.y);
    const result = a.mul(b).normalize().toBytes(.big);
    const expected = stdFeMul(g.x, g.y).toBytes(.big);
    try testing.expectEqualSlices(u8, &expected, &result);
}

test "sq matches stdlib" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const result = a.sq().normalize().toBytes(.big);
    const expected = stdFeSq(g.x).toBytes(.big);
    try testing.expectEqualSlices(u8, &expected, &result);
}

test "sq equals mul(a, a)" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    try testing.expect(a.sq().equivalent(a.mul(a)));

    const b = Fe.fromStdlib(g.y);
    try testing.expect(b.sq().equivalent(b.mul(b)));
}

test "mul(a, invert(a)) == one" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const inv = a.invert();
    const product = a.mul(inv);
    try testing.expect(product.equivalent(Fe.one));
}

test "mul(a, invert(a)) == one for y" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.y);
    const inv = a.invert();
    const product = a.mul(inv);
    try testing.expect(product.equivalent(Fe.one));
}

test "edge: p-1" {
    const p_minus_1 = Fe.fromInt(
        0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2E,
    );
    try testing.expect(!p_minus_1.isZero());
    try testing.expect(p_minus_1.add(Fe.one).isZero());
}

test "fromInt constants" {
    const beta = Fe.fromInt(
        55594575648329892869085402983802832744385952214688224221778511981742606582254,
    );
    try testing.expect(!beta.isZero());
    const beta_cubed = beta.sq().mul(beta);
    try testing.expect(beta_cubed.equivalent(Fe.one));
}

test "dbl matches add(self)" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    try testing.expect(a.dbl().equivalent(a.add(a)));
}

test "mulInt matches repeated add" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    try testing.expect(a.mulInt(3).equivalent(a.add(a).add(a)));
}

test "multiple mul/sq operations stay consistent" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const gx26 = Fe.fromStdlib(g.x);
    const gy26 = Fe.fromStdlib(g.y);

    const result = gx26.sq().mul(gy26);
    const expected = stdFeSq(g.x).mul(g.y);
    try testing.expectEqualSlices(u8, &expected.toBytes(.big), &result.normalize().toBytes(.big));
}

test "normalize idempotent" {
    const g = Secp256k1.basePoint.affineCoordinates();
    const a = Fe.fromStdlib(g.x);
    const n1 = a.normalize();
    const n2 = n1.normalize();
    try testing.expectEqualSlices(u8, &std.mem.toBytes(n1.n), &std.mem.toBytes(n2.n));
}