const std = @import("std");
const crypto = std.crypto;

const IdentityElementError = crypto.errors.IdentityElementError;
const NonCanonicalError = crypto.errors.NonCanonicalError;
const WeakPublicKeyError = crypto.errors.WeakPublicKeyError;

/// Group operations over Curve25519.
pub const Curve25519 = struct {
    /// The underlying prime field.
    pub const Fe = @import("field.zig").Fe;
    /// Field arithmetic mod the order of the main subgroup.
    pub const scalar = @import("scalar.zig");

    x: Fe,

    /// Decode a Curve25519 point from its compressed (X) coordinates.
    pub fn fromBytes(s: [32]u8) Curve25519 {
        return .{ .x = Fe.fromBytes(s) };
    }

    /// Encode a Curve25519 point.
    pub fn toBytes(p: Curve25519) [32]u8 {
        return p.x.toBytes();
    }

    /// The Curve25519 base point.
    pub const basePoint = Curve25519{ .x = Fe.curve25519BasePoint };

    /// Check that the encoding of a Curve25519 point is canonical.
    pub fn rejectNonCanonical(s: [32]u8) NonCanonicalError!void {
        return Fe.rejectNonCanonical(s, false);
    }

    /// Reject the neutral element.
    pub fn rejectIdentity(p: Curve25519) IdentityElementError!void {
        if (p.x.isZero()) {
            return error.IdentityElement;
        }
    }

    /// Multiply a point by the cofactor, returning WeakPublicKey if the element is in a small-order group.
    pub fn clearCofactor(p: Curve25519) WeakPublicKeyError!Curve25519 {
        const cofactor = [_]u8{8} ++ [_]u8{0} ** 31;
        return ladder(p, cofactor, 4) catch return error.WeakPublicKey;
    }

    fn ladder(p: Curve25519, s: [32]u8, comptime bits: usize) IdentityElementError!Curve25519 {
        var x1 = p.x;
        var x2 = Fe.one;
        var z2 = Fe.zero;
        var x3 = x1;
        var z3 = Fe.one;
        var swap: u8 = 0;
        var pos: usize = bits - 1;
        while (true) : (pos -= 1) {
            const bit = (s[pos >> 3] >> @as(u3, @truncate(pos))) & 1;
            swap ^= bit;
            Fe.cSwap2(&x2, &x3, &z2, &z3, swap);
            swap = bit;
            const a = x2.add(z2);
            const b = x2.sub(z2);
            const aa = a.sq();
            const bb = b.sq();
            x2 = aa.mul(bb);
            const e = aa.sub(bb);
            const da = x3.sub(z3).mul(a);
            const cb = x3.add(z3).mul(b);
            x3 = da.add(cb).sq();
            z3 = x1.mul(da.sub(cb).sq());
            z2 = e.mul(bb.add(e.mul32(121666)));
            if (pos == 0) break;
        }
        Fe.cSwap2(&x2, &x3, &z2, &z3, swap);
        z2 = z2.invert();
        x2 = x2.mul(z2);
        if (x2.isZero()) {
            return error.IdentityElement;
        }
        return Curve25519{ .x = x2 };
    }

    /// Multiply a Curve25519 point by a scalar after "clamping" it.
    /// Clamping forces the scalar to be a multiple of the cofactor in
    /// order to prevent small subgroups attacks. This is the standard
    /// way to use Curve25519 for a DH operation.
    /// Return error.IdentityElement if the resulting point is
    /// the identity element.
    pub fn clampedMul(p: Curve25519, s: [32]u8) IdentityElementError!Curve25519 {
        var t: [32]u8 = s;
        scalar.clamp(&t);
        return try ladder(p, t, 255);
    }

    /// Multiply a Curve25519 point by a scalar without clamping it.
    /// Return error.IdentityElement if the resulting point is
    /// the identity element or error.WeakPublicKey if the public
    /// key is a low-order point.
    pub fn mul(p: Curve25519, s: [32]u8) (IdentityElementError || WeakPublicKeyError)!Curve25519 {
        _ = try p.clearCofactor();
        return try ladder(p, s, 256);
    }

    /// Elligator2 map for Curve25519.
    pub fn elligator2(r: Fe) Curve25519 {
        return .{ .x = crypto.ecc.Edwards25519.elligator2(r).x };
    }

    /// Compute the Curve25519 equivalent to an Edwards25519 point.
    ///
    /// Note that the function doesn't check that the input point is
    /// on the prime order group, e.g. that it is an Ed25519 public key
    /// for which an Ed25519 secret key exists.
    ///
    /// If this is required, for example for compatibility with libsodium's strict
    /// validation policy, the caller can call the `rejectUnexpectedSubgroup` function
    /// on the input point before calling this function.
    pub fn fromEdwards25519(p: crypto.ecc.Edwards25519) IdentityElementError!Curve25519 {
        try p.clearCofactor().rejectIdentity();
        const one = crypto.ecc.Edwards25519.Fe.one;
        const py = p.y.mul(p.z.invert());
        const x = one.add(py).mul(one.sub(py).invert()); // xMont=(1+yEd)/(1-yEd)
        return Curve25519{ .x = x };
    }
};

test "curve25519" {
    var s = [32]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8 };
    const p = try Curve25519.basePoint.clampedMul(s);
    try p.rejectIdentity();
    var buf: [128]u8 = undefined;
    try std.testing.expectEqualStrings(try std.fmt.bufPrint(&buf, "{X}", .{&p.toBytes()}), "E6F2A4D1C28EE5C7AD0329268255A468AD407D2672824C0C0EB30EA6EF450145");
    const q = try p.clampedMul(s);
    try std.testing.expectEqualStrings(try std.fmt.bufPrint(&buf, "{X}", .{&q.toBytes()}), "3614E119FFE55EC55B87D6B19971A9F4CBC78EFE80BEC55B96392BABCC712537");

    try Curve25519.rejectNonCanonical(s);
    s[31] |= 0x80;
    try std.testing.expectError(error.NonCanonical, Curve25519.rejectNonCanonical(s));
}

test "non-affine edwards25519 to curve25519 projection" {
    const skh = "90e7595fc89e52fdfddce9c6a43d74dbf6047025ee0462d2d172e8b6a2841d6e";
    var sk: [32]u8 = undefined;
    _ = std.fmt.hexToBytes(&sk, skh) catch unreachable;
    const edp = try crypto.ecc.Edwards25519.basePoint.mul(sk);
    const xp = try Curve25519.fromEdwards25519(edp);
    const expected_hex = "cc4f2cdb695dd766f34118eb67b98652fed1d8bc49c330b119bbfa8a64989378";
    var expected: [32]u8 = undefined;
    _ = std.fmt.hexToBytes(&expected, expected_hex) catch unreachable;
    try std.testing.expectEqualSlices(u8, &xp.toBytes(), &expected);
}

test "elligator2" {
    const Fe = Curve25519.Fe;

    // RFC 9380 Appendix J.4.2 test vector (curve25519_XMD:SHA-512_ELL2_NU_)
    // u = 0x608d892b641f0328523802a6603427c26e55e6f27e71a91a478148d45b5093cd
    // Q.x = 0x51125222da5e763d97f3c10fcc92ea6860b9ccbbd2eb1285728f566721c1e65b
    const u = Fe.fromBytes(.{
        0xcd, 0x93, 0x50, 0x5b, 0xd4, 0x48, 0x81, 0x47, 0x1a, 0xa9, 0x71, 0x7e, 0xf2, 0xe6, 0x55, 0x6e,
        0xc2, 0x27, 0x34, 0x60, 0xa6, 0x02, 0x38, 0x52, 0x28, 0x03, 0x1f, 0x64, 0x2b, 0x89, 0x8d, 0x60,
    });
    const p = Curve25519.elligator2(u);
    try std.testing.expectEqual([32]u8{
        0x5b, 0xe6, 0xc1, 0x21, 0x67, 0x56, 0x8f, 0x72, 0x85, 0x12, 0xeb, 0xd2, 0xbb, 0xcc, 0xb9, 0x60,
        0x68, 0xea, 0x92, 0xcc, 0x0f, 0xc1, 0xf3, 0x97, 0x3d, 0x76, 0x5e, 0xda, 0x22, 0x52, 0x12, 0x51,
    }, p.toBytes());
}

test "small order check" {
    var s: [32]u8 = [_]u8{1} ++ [_]u8{0} ** 31;
    const small_order_ss: [7][32]u8 = .{
        .{
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // 0 (order 4)
        },
        .{
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // 1 (order 1)
        },
        .{
            0xe0, 0xeb, 0x7a, 0x7c, 0x3b, 0x41, 0xb8, 0xae, 0x16, 0x56, 0xe3, 0xfa, 0xf1, 0x9f, 0xc4, 0x6a, 0xda, 0x09, 0x8d, 0xeb, 0x9c, 0x32, 0xb1, 0xfd, 0x86, 0x62, 0x05, 0x16, 0x5f, 0x49, 0xb8, 0x00, // 325606250916557431795983626356110631294008115727848805560023387167927233504 (order 8) */
        },
        .{
            0x5f, 0x9c, 0x95, 0xbc, 0xa3, 0x50, 0x8c, 0x24, 0xb1, 0xd0, 0xb1, 0x55, 0x9c, 0x83, 0xef, 0x5b, 0x04, 0x44, 0x5c, 0xc4, 0x58, 0x1c, 0x8e, 0x86, 0xd8, 0x22, 0x4e, 0xdd, 0xd0, 0x9f, 0x11, 0x57, // 39382357235489614581723060781553021112529911719440698176882885853963445705823 (order 8)
        },
        .{
            0xec, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f, // p-1 (order 2)
        },
        .{
            0xed, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f, // p (=0, order 4)
        },
        .{
            0xee, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f, // p+1 (=1, order 1)
        },
    };
    for (small_order_ss) |small_order_s| {
        try std.testing.expectError(error.WeakPublicKey, Curve25519.fromBytes(small_order_s).clearCofactor());
        try std.testing.expectError(error.WeakPublicKey, Curve25519.fromBytes(small_order_s).mul(s));
        var extra = small_order_s;
        extra[31] ^= 0x80;
        try std.testing.expectError(error.WeakPublicKey, Curve25519.fromBytes(extra).mul(s));
        var valid = small_order_s;
        valid[31] = 0x40;
        s[0] = 0;
        try std.testing.expectError(error.IdentityElement, Curve25519.fromBytes(valid).mul(s));
    }
}
