const std = @import("std");
const Secp256k1 = std.crypto.ecc.Secp256k1;
const Fe = @import("field.zig").Fe;
const StdFe = Secp256k1.Fe;
const StdAffineCoordinates = @TypeOf(Secp256k1.basePoint.affineCoordinates());

/// Affine point in Fe representation.
pub const AffinePoint = struct {
    x: Fe,
    y: Fe,

    pub const identity: AffinePoint = .{ .x = Fe.zero, .y = Fe.zero };

    pub fn neg(self: AffinePoint) AffinePoint {
        return .{ .x = self.x, .y = self.y.neg(1) };
    }

    pub fn fromStdlib(ac: StdAffineCoordinates) AffinePoint {
        return .{
            .x = Fe.fromBytes(ac.x.toBytes(.big), .big),
            .y = Fe.fromBytes(ac.y.toBytes(.big), .big),
        };
    }

    pub fn fromStdlibIdentity(ac: StdAffineCoordinates) AffinePoint {
        // check if this is the identity element (x=0, y=0 in stdlib)
        if (ac.x.isZero() and ac.y.isZero()) return identity;
        return fromStdlib(ac);
    }
};

/// Jacobian point on secp256k1 using Fe: affine (X/Z², Y/Z³).
/// Uses a=0 specialized formulas for fewer field operations than
/// the stdlib's complete projective formulas.
pub const JacobianPoint = struct {
    x: Fe,
    y: Fe,
    z: Fe,

    pub const identity: JacobianPoint = .{ .x = Fe.zero, .y = Fe.one, .z = Fe.zero };

    /// Convert Fe affine (x, y) to Jacobian (x, y, 1).
    pub fn fromAffine(p: AffinePoint) JacobianPoint {
        return .{ .x = p.x, .y = p.y, .z = Fe.one };
    }

    /// Convert stdlib affine to Jacobian via Fe.
    pub fn fromStdlibAffine(ac: StdAffineCoordinates) JacobianPoint {
        return fromAffine(AffinePoint.fromStdlib(ac));
    }

    /// Doubling for a=0 curves (secp256k1). 2M + 5S + 3 normalize.
    /// EFD dbl-2009-l with a=0 specialization.
    pub fn dbl(self: JacobianPoint) JacobianPoint {
        if (self.z.isZero()) return self;

        const a = self.x.sq(); // mag 1
        const b = self.y.sq(); // mag 1
        const c = b.sq(); // mag 1
        const xb = self.x.add(b); // mag 2
        const d = xb.sq().subMag(a, 1).subMag(c, 1).dbl().normalize();
        const e = a.dbl().add(a); // 3a = mag 3
        const f = e.sq();
        const x3 = f.subMag(d.dbl(), 2).normalize();
        const c8 = c.dbl().dbl().dbl(); // mag 8
        const y3 = e.mul(d.subMag(x3, 1)).subMag(c8, 8).normalize();
        const z3 = self.y.mul(self.z).dbl(); // mag 2

        return .{ .x = x3, .y = y3, .z = z3 };
    }

    /// Mixed addition: Jacobian + Affine → Jacobian. 7M + 4S + 4 normalize.
    /// EFD madd-2007-bl. Falls back to dbl() when both points are equal.
    pub fn addMixed(self: JacobianPoint, q: AffinePoint) JacobianPoint {
        if (self.z.isZero()) return fromAffine(q);

        const z1z1 = self.z.sq(); // mag 1
        const qx_z2 = q.x.mul(z1z1); // mag 1
        const s2 = q.y.mul(self.z.mul(z1z1)); // mag 1
        const h = qx_z2.subMag(self.x, 1).normalize();

        // Degenerate case: same x-coordinate
        if (h.isZero()) {
            if (s2.subMag(self.y, 1).normalize().isZero()) return self.dbl();
            return identity;
        }

        const hh = h.sq();
        const i = hh.dbl().dbl(); // mag 4
        const j = h.mul(i);
        const r = s2.subMag(self.y, 1).dbl().normalize();
        const v = self.x.mul(i);
        const x3 = r.sq().subMag(j, 1).subMag(v.dbl(), 2).normalize();
        const y3 = r.mul(v.subMag(x3, 1)).subMag(self.y.mul(j).dbl(), 2).normalize();
        const z3 = self.z.add(h).sq().subMag(z1z1, 1).subMag(hh, 1);

        return .{ .x = x3, .y = y3, .z = z3 };
    }

    /// Mixed subtraction: Jacobian - Affine → Jacobian.
    pub fn subMixed(self: JacobianPoint, q: AffinePoint) JacobianPoint {
        return self.addMixed(q.neg());
    }

    /// Negate the Y coordinate: returns (X, -Y, Z).
    pub fn negY(self: JacobianPoint) JacobianPoint {
        return .{ .x = self.x, .y = self.y.normalize().neg(1), .z = self.z };
    }

    /// Full Jacobian addition: J + J → J. ~12M + 4S.
    /// Needed for combining r1 + r2 at end of verify.
    pub fn add(self: JacobianPoint, other: JacobianPoint) JacobianPoint {
        if (self.z.isZero()) return other;
        if (other.z.isZero()) return self;

        const z1z1 = self.z.sq();
        const z2z2 = other.z.sq();
        const p1 = self.x.mul(z2z2); // X1*Z2²
        const p2 = other.x.mul(z1z1); // X2*Z1²
        const s1 = self.y.mul(other.z.mul(z2z2)); // Y1*Z2³
        const s2 = other.y.mul(self.z.mul(z1z1)); // Y2*Z1³

        // h = p2 - p1 ; both mag 1
        const h = p2.subMag(p1, 1).normalize();
        const r = s2.subMag(s1, 1).normalize();

        // check for doubling case (h == 0 && r == 0)
        if (h.isZero()) {
            if (r.isZero()) return self.dbl();
            return identity;
        }

        const hh = h.sq(); // mag 1
        const hhh = h.mul(hh); // mag 1
        const v = p1.mul(hh); // mag 1
        const r2 = r.sq(); // mag 1
        // x3 = r² - hhh - 2v ; sub(1)=3, sub(2v=2, 2)=6
        const x3 = r2.subMag(hhh, 1).subMag(v.dbl(), 2).normalize();
        // y3 = r*(v-x3) - s1*hhh ; v.sub(x3,1)→3, mul=1; s1.mul(hhh)=1; sub(1)=3
        const y3 = r.mul(v.subMag(x3, 1)).subMag(s1.mul(hhh), 1).normalize();
        const z3 = self.z.mul(other.z).mul(h); // mag 1

        return .{ .x = x3, .y = y3, .z = z3 };
    }

    /// Compare signature r against R's x-coordinate in Jacobian space.
    /// Checks r * Z² == X (mod p), avoiding field inversion.
    /// Returns true if the x-coordinates match.
    pub fn jacobianCompare(self: JacobianPoint, r_fe: Fe) bool {
        // Jacobian: x_affine = X / Z²
        // So check: r * Z² == X
        const z2 = self.z.sq();
        const rz2 = r_fe.mul(z2);
        if (rz2.equivalent(self.x)) return true;

        // rare overflow case: x_affine could be in [n, p)
        // check (r + n) * Z² == X
        const n_fe26 = comptime Fe.fromInt(Secp256k1.scalar.field_order);
        const p_minus_n: u256 = StdFe.field_order - Secp256k1.scalar.field_order;
        const r_norm = r_fe.normalize();
        const r_bytes = r_norm.toBytes(.big);
        const r_int = std.mem.readInt(u256, &r_bytes, .big);
        if (r_int < p_minus_n) {
            const rn_z2 = r_fe.add(n_fe26).mul(z2);
            if (rn_z2.equivalent(self.x)) return true;
        }

        return false;
    }

    /// Convert to stdlib projective for backward compatibility.
    /// Jacobian: x = X/Z², y = Y/Z³. Projective: x = X'/Z', y = Y'/Z'.
    /// Set X' = X*Z, Y' = Y, Z' = Z³.
    pub fn toProjective(self: JacobianPoint) Secp256k1 {
        if (self.z.isZero()) return Secp256k1.identityElement;
        return .{
            .x = self.x.mul(self.z).toStdlib(),
            .y = self.y.toStdlib(),
            .z = self.z.sq().mul(self.z).toStdlib(),
        };
    }
};

// ============================================================
// tests — verify Fe Jacobian matches stdlib results
// ============================================================

fn toStdlibAffine(j: JacobianPoint) StdAffineCoordinates {
    return j.toProjective().affineCoordinates();
}

test "dbl matches stdlib" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();

    const j_2g = JacobianPoint.fromStdlibAffine(g_affine).dbl();
    const j_2g_affine = toStdlibAffine(j_2g);

    const stdlib_2g_affine = G.dbl().affineCoordinates();

    try std.testing.expect(j_2g_affine.x.equivalent(stdlib_2g_affine.x));
    try std.testing.expect(j_2g_affine.y.equivalent(stdlib_2g_affine.y));
}

test "addMixed matches stdlib add" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    // 2G + G = 3G
    const j_3g = JacobianPoint.fromAffine(g26).dbl().addMixed(g26);
    const j_3g_affine = toStdlibAffine(j_3g);

    const stdlib_3g_affine = G.dbl().add(G).affineCoordinates();

    try std.testing.expect(j_3g_affine.x.equivalent(stdlib_3g_affine.x));
    try std.testing.expect(j_3g_affine.y.equivalent(stdlib_3g_affine.y));
}

test "identity + affine = affine" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    const result = JacobianPoint.identity.addMixed(g26);
    const result_affine = toStdlibAffine(result);

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

test "subMixed is inverse of addMixed" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    // 3G - G = 2G
    const j_3g = JacobianPoint.fromAffine(g26).dbl().addMixed(g26);
    const j_2g = j_3g.subMixed(g26);
    const j_2g_affine = toStdlibAffine(j_2g);

    const stdlib_2g_affine = G.dbl().affineCoordinates();

    try std.testing.expect(j_2g_affine.x.equivalent(stdlib_2g_affine.x));
    try std.testing.expect(j_2g_affine.y.equivalent(stdlib_2g_affine.y));
}

test "repeated doubling" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();

    // 16G via 4 doublings
    const j = JacobianPoint.fromStdlibAffine(g_affine).dbl().dbl().dbl().dbl();
    const j_affine = toStdlibAffine(j);

    const stdlib_affine = G.dbl().dbl().dbl().dbl().affineCoordinates();

    try std.testing.expect(j_affine.x.equivalent(stdlib_affine.x));
    try std.testing.expect(j_affine.y.equivalent(stdlib_affine.y));
}

test "full Jacobian add matches stdlib" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    // 2G + 3G = 5G
    const j_2g = JacobianPoint.fromAffine(g26).dbl();
    const j_3g = j_2g.addMixed(g26);
    const j_5g = j_2g.add(j_3g);
    const j_5g_affine = toStdlibAffine(j_5g);

    const stdlib_5g = G.dbl().add(G.dbl().add(G));
    const stdlib_5g_affine = stdlib_5g.affineCoordinates();

    try std.testing.expect(j_5g_affine.x.equivalent(stdlib_5g_affine.x));
    try std.testing.expect(j_5g_affine.y.equivalent(stdlib_5g_affine.y));
}

test "add with identity" {
    const G = Secp256k1.basePoint;
    const g_affine = G.affineCoordinates();
    const g26 = AffinePoint.fromStdlib(g_affine);

    const j_g = JacobianPoint.fromAffine(g26);
    const result = j_g.add(JacobianPoint.identity);
    const result_affine = toStdlibAffine(result);

    try std.testing.expect(result_affine.x.equivalent(g_affine.x));
    try std.testing.expect(result_affine.y.equivalent(g_affine.y));

    const result2 = JacobianPoint.identity.add(j_g);
    const result2_affine = toStdlibAffine(result2);

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