const std = @import("std");
const Secp256k1 = std.crypto.ecc.Secp256k1;
const Fe = @import("field.zig").Fe;
const scalar = Secp256k1.scalar;

const endo = @import("endo.zig");
const point = @import("point.zig");
const affine_mod = @import("affine.zig");
const jacobian_mod = @import("jacobian.zig");
const JacobianPoint = jacobian_mod.JacobianPoint;
const AffinePoint = jacobian_mod.AffinePoint;

pub const VerifyError = std.crypto.errors.IdentityElementError ||
    std.crypto.errors.NonCanonicalError ||
    error{SignatureVerificationFailed};

/// Precomputed base point table: gTable()[i][j] = j * 256^i * G in affine.
/// 32 subtables × 256 entries = 8192 affine points.
/// Enables u1*G computation via byte-at-a-time lookup with zero doublings.
/// Lazily initialized at runtime (u128 field arithmetic is too heavy
/// for the comptime interpreter at this scale).
var g_table_storage: [32][256]AffinePoint = undefined;
var g_table_ready: bool = false;

fn gTable() *const [32][256]AffinePoint {
    if (!g_table_ready) {
        g_table_storage = affine_mod.buildByteTable(Secp256k1.basePoint);
        g_table_ready = true;
    }
    return &g_table_storage;
}

/// Reduce a 32-byte big-endian value to a scalar (mod n).
fn reduceToScalar(h: [32]u8) scalar.Scalar {
    var xs = [_]u8{0} ** 48;
    @memcpy(xs[xs.len - 32 ..], &h);
    return scalar.Scalar.fromBytes48(xs, .big);
}

/// Verify an ECDSA signature using precomputed tables and Jacobian arithmetic.
///
/// Optimizations:
/// 1. u1*G via precomputed byte table: ~32 mixed adds, zero doublings
/// 2. u2*Q via projective-table GLV: no field inversion, 4-bit windowed
/// 3. Jacobian comparison (no field inversion for final check)
/// 4. All field arithmetic via 5×52-bit limbs (libsecp256k1 style) — fast on 64-bit hardware
pub fn verify(sig_r: [32]u8, sig_s: [32]u8, msg_hash: [32]u8, public_key: Secp256k1) VerifyError!void {
    // parse and validate r, s
    const r_sc = scalar.Scalar.fromBytes(sig_r, .big) catch return error.SignatureVerificationFailed;
    const s_sc = scalar.Scalar.fromBytes(sig_s, .big) catch return error.SignatureVerificationFailed;
    if (r_sc.isZero() or s_sc.isZero()) return error.IdentityElement;

    // scalar_u1 = z * s^-1, scalar_u2 = r * s^-1
    const z = reduceToScalar(msg_hash);
    const s_inv = scalarInvert(s_sc);
    const scalar_u1 = z.mul(s_inv).toBytes(.big);
    const scalar_u2 = r_sc.mul(s_inv).toBytes(.little);

    // 1. u1*G via precomputed byte table — direct 32-byte lookup, no GLV
    const r1 = basePointMul(scalar_u1);

    // 2. u2*Q via projective-table GLV — no field inversion
    var split_u2 = endo.splitScalar(scalar_u2, .little) catch return error.SignatureVerificationFailed;
    const zero_s = scalar.Scalar.zero.toBytes(.little);
    var neg_p = false;
    var neg_p_phi = false;
    if (split_u2.r1[16] != 0) {
        split_u2.r1 = scalar.neg(split_u2.r1, .little) catch zero_s;
        neg_p = true;
    }
    if (split_u2.r2[16] != 0) {
        split_u2.r2 = scalar.neg(split_u2.r2, .little) catch zero_s;
        neg_p_phi = true;
    }
    const pk_affine26 = AffinePoint.fromStdlib(public_key.affineCoordinates());
    const r2 = publicKeyMulProjective(split_u2, neg_p, neg_p_phi, pk_affine26);

    // 3. combine results in Jacobian, compare without inversion
    const q = r1.add(r2);

    // reject identity
    if (q.z.isZero()) return error.SignatureVerificationFailed;

    // Jacobian comparison: r * Z² == X (mod p)
    const r_fe = Fe.fromBytes(sig_r, .big);
    if (!q.jacobianCompare(r_fe)) {
        return error.SignatureVerificationFailed;
    }
}

/// Compute u1*G using the precomputed byte table.
/// Direct 32-byte lookup: G_TABLE[i][scalar[i]], no GLV decomposition.
/// Cost: ~32 mixed Jacobian-affine additions, zero doublings.
fn basePointMul(scalar_u1: [32]u8) JacobianPoint {
    const table = gTable();
    var acc = JacobianPoint.identity;

    // table[i] corresponds to byte position i (little-endian).
    // scalar_u1 is big-endian, so byte 31 is least significant.
    // table[i] ↔ scalar_u1[31-i]
    for (0..32) |i| {
        const byte = scalar_u1[31 - i];
        if (byte != 0) {
            acc = acc.addMixed(table[i][byte]);
        }
    }

    return acc;
}

/// Compute u2*Q using 2-way Jacobian Shamir with projective tables.
/// Tables stay in Jacobian form — no batchToAffine field inversion.
/// Cost: 128 Jacobian doublings + ~44 Jacobian additions.
fn publicKeyMulProjective(
    split: endo.SplitScalar,
    neg_p: bool,
    neg_p_phi: bool,
    pk_affine: AffinePoint,
) JacobianPoint {
    // Build 9-entry Jacobian tables: [identity, 1P, 2P, ..., 8P]
    const pk_table = point.precompute(pk_affine, 8);
    const pk_phi_table = point.phiTableJacobian(9, pk_table);

    const e1 = point.slide(split.r1);
    const e2 = point.slide(split.r2);

    var q = JacobianPoint.identity;
    var pos: usize = 32; // 128-bit half-scalars → 32 nybbles + carry
    while (true) : (pos -= 1) {
        q = addSlotJ(q, &pk_table, e1[pos], neg_p);
        q = addSlotJ(q, &pk_phi_table, e2[pos], neg_p_phi);
        if (pos == 0) break;
        q = q.dbl().dbl().dbl().dbl();
    }
    return q;
}

/// Add/subtract a Jacobian table entry based on signed digit and negation flag.
/// Variable-time: branches on digit value (safe for public verification).
inline fn addSlotJ(q: JacobianPoint, table: *const [9]JacobianPoint, slot: i8, negate: bool) JacobianPoint {
    var s = slot;
    if (negate) s = -s;
    if (s > 0) {
        return q.add(table[@intCast(s)]);
    } else if (s < 0) {
        return q.add(table[@intCast(-s)].negY());
    }
    return q;
}

/// Fermat scalar inversion: s^(n-2) mod n via addition chain.
/// 253 squarings + 40 multiplications. Variable-time (safe for verification
/// where all inputs are public).
///
/// Addition chain generated by github.com/mmcloughlin/addchain v0.4.0,
/// ported from gitlab.com/yawning/secp256k1-voi.
fn scalarInvert(x: scalar.Scalar) scalar.Scalar {
    // Step 1: t0 = x^0x2
    var t0 = x.sq();

    // Step 2: t1 = x^0x3
    var t1 = x.mul(t0);

    // Step 3: t2 = x^0x5
    var t2 = t0.mul(t1);

    // Step 4: t3 = x^0x7
    var t3 = t0.mul(t2);

    // Step 5: t4 = x^0x9
    var t4 = t0.mul(t3);

    // Step 6: t5 = x^0xb
    var t5 = t0.mul(t4);

    // Step 7: t0 = x^0xd
    t0 = t0.mul(t5);

    // Step 9: t6 = x^0x34
    var t6 = pow2k(t0, 2);

    // Step 10: t6 = x^0x3f
    t6 = t5.mul(t6);

    // Step 11: t7 = x^0x7e
    var t7 = t6.sq();

    // Step 12: t7 = x^0x7f
    t7 = x.mul(t7);

    // Step 13: t8 = x^0xfe
    var t8 = t7.sq();

    // Step 14: t8 = x^0xff
    t8 = x.mul(t8);

    // Step 17: t9 = x^0x7f8
    var t9 = pow2k(t8, 3);

    // Step 19: t10 = x^0x1fe0
    var t10 = pow2k(t9, 2);

    // Step 20: t11 = x^0x3fc0
    var t11 = t10.sq();

    // Step 21: t12 = x^0x7f80
    var t12 = t11.sq();

    // Step 28: t13 = x^0x3fc000
    const t13 = pow2k(t12, 7);

    // Step 29: t11 = x^0x3fffc0
    t11 = t11.mul(t13);

    // Step 38: t11 = x^0x7fff8000
    t11 = pow2k(t11, 9);

    // Step 39: t12 = x^0x7fffff80
    t12 = t12.mul(t11);

    // Step 45: t11 = x^0x1fffffe000
    t11 = pow2k(t12, 6);

    // Step 46: t10 = x^0x1fffffffe0
    t10 = t10.mul(t11);

    // Step 72: t10 = x^0x7fffffff80000000
    t10 = pow2k(t10, 26);

    // Step 73: t12 = x^0x7fffffffffffff80
    t12 = t12.mul(t10);

    // Step 77: t10 = x^0x7fffffffffffff800
    t10 = pow2k(t12, 4);

    // Step 78: t9 = x^0x7fffffffffffffff8
    t9 = t9.mul(t10);

    // Step 138: t9 = x^0x7fffffffffffffff8000000000000000
    t9 = pow2k(t9, 60);

    // Step 139: t12 = x^0x7fffffffffffffffffffffffffffff80
    t12 = t12.mul(t9);

    // Step 140: t7 = x^0x7fffffffffffffffffffffffffffffff
    t7 = t7.mul(t12);

    // Step 145: t7 = x^0xfffffffffffffffffffffffffffffffe0
    t7 = pow2k(t7, 5);

    // Step 146: t7 = x^0xfffffffffffffffffffffffffffffffeb
    t7 = t5.mul(t7);

    // Step 149: t7 = x^0x7fffffffffffffffffffffffffffffff58
    t7 = pow2k(t7, 3);

    // Step 150: t7 = x^0x7fffffffffffffffffffffffffffffff5d
    t7 = t2.mul(t7);

    // Step 154: t7 = x^0x7fffffffffffffffffffffffffffffff5d0
    t7 = pow2k(t7, 4);

    // Step 155: t7 = x^0x7fffffffffffffffffffffffffffffff5d5
    t7 = t2.mul(t7);

    // Step 159: t7 = x^0x7fffffffffffffffffffffffffffffff5d50
    t7 = pow2k(t7, 4);

    // Step 160: t7 = x^0x7fffffffffffffffffffffffffffffff5d57
    t7 = t3.mul(t7);

    // Step 165: t7 = x^0xfffffffffffffffffffffffffffffffebaae0
    t7 = pow2k(t7, 5);

    // Step 166: t7 = x^0xfffffffffffffffffffffffffffffffebaaed
    t7 = t0.mul(t7);

    // Step 168: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb4
    t7 = pow2k(t7, 2);

    // Step 169: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb7
    t7 = t1.mul(t7);

    // Step 174: t7 = x^0x7fffffffffffffffffffffffffffffff5d576e0
    t7 = pow2k(t7, 5);

    // Step 175: t7 = x^0x7fffffffffffffffffffffffffffffff5d576e7
    t7 = t3.mul(t7);

    // Step 181: t7 = x^0x1fffffffffffffffffffffffffffffffd755db9c0
    t7 = pow2k(t7, 6);

    // Step 182: t7 = x^0x1fffffffffffffffffffffffffffffffd755db9cd
    t7 = t0.mul(t7);

    // Step 187: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb739a0
    t7 = pow2k(t7, 5);

    // Step 188: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb739ab
    t7 = t5.mul(t7);

    // Step 192: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb739ab0
    t7 = pow2k(t7, 4);

    // Step 193: t7 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd
    t7 = t0.mul(t7);

    // Step 196: t7 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e8
    t7 = pow2k(t7, 3);

    // Step 197: t7 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e9
    t7 = x.mul(t7);

    // Step 203: t7 = x^0x7fffffffffffffffffffffffffffffff5d576e7357a40
    t7 = pow2k(t7, 6);

    // Step 204: t2 = x^0x7fffffffffffffffffffffffffffffff5d576e7357a45
    t2 = t2.mul(t7);

    // Step 214: t2 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e91400
    t2 = pow2k(t2, 10);

    // Step 215: t2 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e91407
    t2 = t3.mul(t2);

    // Step 219: t2 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e914070
    t2 = pow2k(t2, 4);

    // Step 220: t3 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e914077
    t3 = t3.mul(t2);

    // Step 229: t3 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280ee00
    t3 = pow2k(t3, 9);

    // Step 230: t8 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff
    t8 = t8.mul(t3);

    // Step 235: t8 = x^0x7fffffffffffffffffffffffffffffff5d576e7357a4501ddfe0
    t8 = pow2k(t8, 5);

    // Step 236: t8 = x^0x7fffffffffffffffffffffffffffffff5d576e7357a4501ddfe9
    t8 = t4.mul(t8);

    // Step 242: t8 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e9140777fa40
    t8 = pow2k(t8, 6);

    // Step 243: t5 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e9140777fa4b
    t5 = t5.mul(t8);

    // Step 247: t5 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e9140777fa4b0
    t5 = pow2k(t5, 4);

    // Step 248: t5 = x^0x1fffffffffffffffffffffffffffffffd755db9cd5e9140777fa4bd
    t5 = t0.mul(t5);

    // Step 253: t5 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a0
    t5 = pow2k(t5, 5);

    // Step 254: t1 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a3
    t1 = t1.mul(t5);

    // Step 260: t1 = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8c0
    t1 = pow2k(t1, 6);

    // Step 261: t1 = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd
    t1 = t0.mul(t1);

    // Step 271: t1 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a33400
    t1 = pow2k(t1, 10);

    // Step 272: t0 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a3340d
    t0 = t0.mul(t1);

    // Step 276: t0 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a3340d0
    t0 = pow2k(t0, 4);

    // Step 277: t4 = x^0x3fffffffffffffffffffffffffffffffaeabb739abd2280eeff497a3340d9
    t4 = t4.mul(t0);

    // Step 283: t4 = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd03640
    t4 = pow2k(t4, 6);

    // Step 284: t14 = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd03641
    var t14 = x.mul(t4);

    // Step 292: t14 = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364100
    t14 = pow2k(t14, 8);

    // Step 293: result = x^0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd036413f
    // This is x^(n-2) where n is the secp256k1 group order
    return t6.mul(t14);
}

/// k repeated squarings: s^(2^k)
inline fn pow2k(s: scalar.Scalar, comptime k: comptime_int) scalar.Scalar {
    var r = s;
    inline for (0..k) |_| {
        r = r.sq();
    }
    return r;
}

test "scalarInvert matches stdlib invert" {
    const S = scalar.Scalar;

    // test with scalar = 1
    const one = S.one;
    const inv_one = scalarInvert(one);
    const inv_one_std = one.invert();
    try std.testing.expectEqual(inv_one.toBytes(.big), inv_one_std.toBytes(.big));

    // test with a known scalar from base point
    const s_bytes = Secp256k1.basePoint.affineCoordinates().x.toBytes(.big);
    const s_val = S.fromBytes(s_bytes, .big) catch unreachable;
    const inv_fermat = scalarInvert(s_val);
    const inv_stdlib = s_val.invert();
    try std.testing.expectEqual(inv_fermat.toBytes(.big), inv_stdlib.toBytes(.big));

    // verify: s * s^-1 == 1
    const product = s_val.mul(inv_fermat);
    try std.testing.expectEqual(product.toBytes(.big), one.toBytes(.big));

    // test with n-1 (the largest valid scalar)
    // n-1 in big-endian
    const n_minus_1_bytes = [_]u8{
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
        0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFE,
        0xBA, 0xAE, 0xDC, 0xE6, 0xAF, 0x48, 0xA0, 0x3B,
        0xBF, 0xD2, 0x5E, 0x8C, 0xD0, 0x36, 0x41, 0x40,
    };
    const nm1 = S.fromBytes(n_minus_1_bytes, .big) catch unreachable;
    const inv_nm1_fermat = scalarInvert(nm1);
    const inv_nm1_stdlib = nm1.invert();
    try std.testing.expectEqual(inv_nm1_fermat.toBytes(.big), inv_nm1_stdlib.toBytes(.big));

    // test with several more deterministic values
    inline for ([_]u256{ 2, 7, 42, 0xDEADBEEF, 0x123456789ABCDEF0 }) |v| {
        var buf: [32]u8 = undefined;
        std.mem.writeInt(u256, &buf, v, .big);
        const sv = S.fromBytes(buf, .big) catch unreachable;
        const inv_f = scalarInvert(sv);
        const inv_s = sv.invert();
        try std.testing.expectEqual(inv_f.toBytes(.big), inv_s.toBytes(.big));
        // verify round-trip
        try std.testing.expectEqual(sv.mul(inv_f).toBytes(.big), one.toBytes(.big));
    }
}

test "gTable()[0][1] is the base point" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const table = gTable();

    const table_g = table[0][1];
    const table_g_x = table_g.x.toStdlib();
    const table_g_y = table_g.y.toStdlib();

    try std.testing.expect(table_g_x.equivalent(g_affine.x));
    try std.testing.expect(table_g_y.equivalent(g_affine.y));
}

test "gTable()[0][2] is 2G" {
    const G = Secp256k1.basePoint;
    const g2 = G.dbl().affineCoordinates();
    const table = gTable();

    const table_2g = table[0][2];
    const table_2g_x = table_2g.x.toStdlib();
    const table_2g_y = table_2g.y.toStdlib();

    try std.testing.expect(table_2g_x.equivalent(g2.x));
    try std.testing.expect(table_2g_y.equivalent(g2.y));
}

test "comptime addMixed gives correct 2G" {
    @setEvalBranchQuota(100_000_000);
    const g_affine = comptime AffinePoint.fromStdlib(Secp256k1.basePoint.affineCoordinates());
    const j_g = comptime JacobianPoint.fromAffine(g_affine);
    const j_2g = comptime j_g.addMixed(g_affine);

    // convert to affine via toProjective
    const proj = j_2g.toProjective();
    const affine = proj.affineCoordinates();

    const expected = Secp256k1.basePoint.dbl().affineCoordinates();

    try std.testing.expect(affine.x.equivalent(expected.x));
    try std.testing.expect(affine.y.equivalent(expected.y));
}

test "projective x-coordinate: X == x_affine * Z" {
    // verify the coordinate convention: x_affine = X / Z (projective, not Jacobian)
    const s = comptime blk: {
        var buf: [32]u8 = undefined;
        std.mem.writeInt(u256, &buf, 12345, .little);
        break :blk buf;
    };
    const P = try Secp256k1.basePoint.mul(s, .little);
    const affine = P.affineCoordinates();
    const x_via_z = affine.x.mul(P.z);
    try std.testing.expect(P.x.equivalent(x_via_z));
}