const std = @import("std");
const jacobian_mod = @import("jacobian.zig");
const JacobianPoint = jacobian_mod.JacobianPoint;
const AffinePoint = jacobian_mod.AffinePoint;
const endo = @import("endo.zig");
const affine_mod = @import("affine.zig");
const Secp256k1 = std.crypto.ecc.Secp256k1;

/// Build a precomputation table: [identity, p, 2p, 3p, ..., count*p].
/// Uses Jacobian arithmetic for efficient computation.
pub fn precompute(p: AffinePoint, comptime count: usize) [1 + count]JacobianPoint {
    var pc: [1 + count]JacobianPoint = undefined;
    pc[0] = JacobianPoint.identity;
    pc[1] = JacobianPoint.fromAffine(p);
    var i: usize = 2;
    while (i <= count) : (i += 1) {
        pc[i] = if (i % 2 == 0) pc[i / 2].dbl() else pc[i - 1].addMixed(p);
    }
    return pc;
}

/// Apply the endomorphism to each entry in an affine precompute table.
/// phi(n*P) = n*phi(P), so transforming the table gives a valid table for phi(P).
pub fn phiTableAffine(comptime n: usize, table: [n]AffinePoint) [n]AffinePoint {
    var result: [n]AffinePoint = undefined;
    for (0..n) |i| {
        result[i] = endo.phiAffine(table[i]);
    }
    return result;
}

/// Apply the endomorphism to each entry in a Jacobian precompute table.
pub fn phiTableJacobian(comptime n: usize, table: [n]JacobianPoint) [n]JacobianPoint {
    var result: [n]JacobianPoint = undefined;
    for (0..n) |i| {
        result[i] = endo.phiJacobian(table[i]);
    }
    return result;
}

/// Encode a scalar into a signed 4-bit windowed representation.
/// Returns 65 digits in [-8, 8], suitable for use with a 9-entry precompute table.
/// For half-sized scalars (128-bit), digits 33-64 are zero.
pub fn slide(s: [32]u8) [2 * 32 + 1]i8 {
    var e: [2 * 32 + 1]i8 = undefined;
    for (s, 0..) |x, i| {
        e[i * 2 + 0] = @as(i8, @as(u4, @truncate(x)));
        e[i * 2 + 1] = @as(i8, @as(u4, @truncate(x >> 4)));
    }
    var carry: i8 = 0;
    for (e[0..64]) |*x| {
        x.* += carry;
        carry = (x.* + 8) >> 4;
        x.* -= carry * 16;
    }
    e[64] = carry;
    return e;
}

test "precompute table correctness" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    const table = precompute(g26, 8);

    // table[0] should be identity
    try std.testing.expect(table[0].z.isZero());

    // table[1] should be G — compare via toProjective
    const t1_affine = table[1].toProjective().affineCoordinates();
    try std.testing.expect(t1_affine.x.equivalent(g_affine.x));
    try std.testing.expect(t1_affine.y.equivalent(g_affine.y));

    // table[2] should be 2G
    const t2_affine = table[2].toProjective().affineCoordinates();
    const stdlib_2g = G.dbl().affineCoordinates();
    try std.testing.expect(t2_affine.x.equivalent(stdlib_2g.x));
    try std.testing.expect(t2_affine.y.equivalent(stdlib_2g.y));

    // table[3] should be 3G = 2G + G
    const t3_affine = table[3].toProjective().affineCoordinates();
    const stdlib_3g = G.dbl().add(G).affineCoordinates();
    try std.testing.expect(t3_affine.x.equivalent(stdlib_3g.x));
    try std.testing.expect(t3_affine.y.equivalent(stdlib_3g.y));
}

test "slide produces valid digits" {
    // Test with a known scalar
    var s: [32]u8 = undefined;
    std.mem.writeInt(u256, &s, 12345678, .little);
    const digits = slide(s);

    // All digits should be in [-8, 8]
    for (digits) |d| {
        try std.testing.expect(d >= -8 and d <= 8);
    }

    // Reconstruct the value and verify
    var reconstructed: i512 = 0;
    var power: i512 = 1;
    for (digits) |d| {
        reconstructed += @as(i512, d) * power;
        power *= 16;
    }
    try std.testing.expectEqual(@as(i512, 12345678), reconstructed);
}