const std = @import("std");
const Secp256k1 = std.crypto.ecc.Secp256k1;
const Fe = @import("field.zig").Fe;
const jacobian_mod = @import("jacobian.zig");
const JacobianPoint = jacobian_mod.JacobianPoint;
const AffinePoint = jacobian_mod.AffinePoint;

/// Batch-convert Jacobian points to affine using Montgomery's trick.
/// Uses 1 field inversion + 3*(n-1) field multiplications.
/// For Jacobian: x_affine = X * Z^{-2}, y_affine = Y * Z^{-3}.
pub fn batchToAffine(comptime n: usize, points: [n]JacobianPoint) [n]AffinePoint {
    // Replace zero Z with one to keep products invertible
    var zs: [n]Fe = undefined;
    for (0..n) |i| {
        zs[i] = if (points[i].z.isZero()) Fe.one else points[i].z;
    }

    // Forward: accumulate Z products
    var products: [n]Fe = undefined;
    products[0] = zs[0];
    for (1..n) |i| {
        products[i] = products[i - 1].mul(zs[i]);
    }

    // Invert total product
    var inv = products[n - 1].invert();

    // Backward: recover individual Z inverses and convert to affine
    // For Jacobian: x = X * zinv², y = Y * zinv³
    var result: [n]AffinePoint = undefined;
    var i: usize = n - 1;
    while (true) {
        const z_inv = if (i > 0) inv.mul(products[i - 1]) else inv;
        if (i > 0) inv = inv.mul(zs[i]);

        if (points[i].z.isZero()) {
            result[i] = AffinePoint.identity;
        } else {
            const z_inv2 = z_inv.sq();
            const z_inv3 = z_inv2.mul(z_inv);
            result[i] = .{
                .x = points[i].x.mul(z_inv2),
                .y = points[i].y.mul(z_inv3),
            };
        }

        if (i == 0) break;
        i -= 1;
    }

    return result;
}

/// Build a byte-indexed precomputed table for scalar multiplication.
/// table[i][j] = j * 256^i * base, stored as AffinePoint.
/// 32 subtables × 256 entries = full 256-bit scalar coverage.
/// table[i][0] is the identity element (unused in lookups).
pub fn buildByteTable(base: Secp256k1) [32][256]AffinePoint {
    @setEvalBranchQuota(100_000_000);

    // Phase 1: compute all 32*256 points in Jacobian Fe
    var flat: [32 * 256]JacobianPoint = undefined;

    var cur_base_affine = AffinePoint.fromStdlib(base.affineCoordinates());
    for (0..32) |sub| {
        flat[sub * 256] = JacobianPoint.identity;
        flat[sub * 256 + 1] = JacobianPoint.fromAffine(cur_base_affine);
        for (2..256) |j| {
            flat[sub * 256 + j] = flat[sub * 256 + j - 1].addMixed(cur_base_affine);
        }
        if (sub < 31) {
            // next subtable base = 256 * current base (8 doublings)
            var next = JacobianPoint.fromAffine(cur_base_affine);
            next = next.dbl().dbl().dbl().dbl().dbl().dbl().dbl().dbl();
            // convert back to affine for next iteration
            const next_batch = batchToAffine(1, .{next});
            cur_base_affine = next_batch[0];
        }
    }

    // Phase 2: batch convert to affine
    const affine_flat = batchToAffine(32 * 256, flat);

    // Reshape to [32][256]
    var result: [32][256]AffinePoint = undefined;
    for (0..32) |sub| {
        for (0..256) |j| {
            result[sub][j] = affine_flat[sub * 256 + j];
        }
    }
    return result;
}

test "batchToAffine matches individual conversion" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);
    const StdAffineCoordinates = @TypeOf(g_affine);

    const j_id = JacobianPoint.identity;
    const j_g = JacobianPoint.fromAffine(g26);
    const j_2g = j_g.dbl();
    const j_3g = j_2g.addMixed(g26);

    const points = [_]JacobianPoint{ j_id, j_g, j_2g, j_3g };
    const result = batchToAffine(4, points);

    // Identity
    try std.testing.expect(result[0].x.isZero());

    // G — compare via stdlib
    const result_g = StdAffineCoordinates{
        .x = result[1].x.toStdlib(),
        .y = result[1].y.toStdlib(),
    };
    try std.testing.expect(result_g.x.equivalent(g_affine.x));
    try std.testing.expect(result_g.y.equivalent(g_affine.y));

    // 2G
    const g2_affine = G.dbl().affineCoordinates();
    const result_2g = StdAffineCoordinates{
        .x = result[2].x.toStdlib(),
        .y = result[2].y.toStdlib(),
    };
    try std.testing.expect(result_2g.x.equivalent(g2_affine.x));
    try std.testing.expect(result_2g.y.equivalent(g2_affine.y));

    // 3G
    const g3_affine = G.dbl().add(G).affineCoordinates();
    const result_3g = StdAffineCoordinates{
        .x = result[3].x.toStdlib(),
        .y = result[3].y.toStdlib(),
    };
    try std.testing.expect(result_3g.x.equivalent(g3_affine.x));
    try std.testing.expect(result_3g.y.equivalent(g3_affine.y));
}