import gmpy2
from gmpy2 import mpz, is_prime, powmod, invert, random_state, mpz_random
import random

# Task 1: Generate a random elliptic curve over F_p
def generate_random_curve(p):
    p = mpz(p)
    rs = random_state(random.randrange(2**32))
    while True:
        A = mpz_random(rs, p)
        B = mpz_random(rs, p)
        # Check that the discriminant is non-zero
        discriminant = (4 * powmod(A, 3, p) + 27 * powmod(B, 2, p)) % p
        if discriminant != 0:
            return A, B

# Task 2: Find a random point on the elliptic curve over F_p
def find_random_point(A, B, p):
    p = mpz(p)
    rs = random_state(random.randrange(2**32))
    while True:
        x = mpz_random(rs, p)
        rhs = (powmod(x, 3, p) + A * x + B) % p
        legendre_symbol = gmpy2.legendre(rhs, p)
        if legendre_symbol == 1:
            # Compute y using the fact that p ≡ 3 mod 4
            y = powmod(rhs, (p + 1) // 4, p)
            return x, y
        elif legendre_symbol == 0:
            # y = 0 is a valid solution
            y = mpz(0)
            return x, y
        # If rhs is not a quadratic residue, try another x

# Task 3: Compute the inverse (negative) of a given point
def point_negation(P, p):
    x, y = P
    return x, (-y) % p

# Task 4: Compute P ⊕ Q, the sum of two points on the elliptic curve
def point_addition(P, Q, A, p):
    if P == 'O':
        return Q
    if Q == 'O':
        return P
    x1, y1 = P
    x2, y2 = Q
    if x1 == x2 and (y1 + y2) % p == 0:
        # P + (-P) = O
        return 'O'
    if x1 == x2 and y1 == y2:
        # Point doubling
        if y1 == 0:
            return 'O'
        numerator = (3 * powmod(x1, 2, p) + A) % p
        denominator = (2 * y1) % p
    else:
        # P ≠ Q
        numerator = (y2 - y1) % p
        denominator = (x2 - x1) % p
    inv_denominator = invert(denominator, p)
    if inv_denominator is None:
        return 'O'  # Cannot invert denominator; result is point at infinity
    lam = (numerator * inv_denominator) % p
    x3 = (powmod(lam, 2, p) - x1 - x2) % p
    y3 = (lam * (x1 - x3) - y1) % p
    return x3, y3

# Task 5: Compute nP, the scalar multiplication of a point
def scalar_multiplication(P, n, A, p):
    n = mpz(n)
    result = 'O'  # Start with the point at infinity
    addend = P
    while n > 0:
        if n % 2 == 1:
            result = point_addition(result, addend, A, p)
        addend = point_addition(addend, addend, A, p)
        n = n // 2
    return result

# Helper function to generate a large prime p ≡ 3 mod 4
def generate_large_prime(bits):
    rs = random_state(random.randrange(2**32))
    while True:
        candidate = mpz_random(rs, 2 ** bits)
        candidate |= 1  # Ensure it's odd
        candidate |= 3  # Ensure p ≡ 3 mod 4
        candidate = gmpy2.next_prime(candidate)
        if candidate % 4 == 3 and candidate.bit_length() == bits:
            return candidate

# Example usage of the functions
if __name__ == "__main__":
    # Generate a large prime p ≡ 3 mod 4 (approximately 300 bits)
    bits = 300
    p = generate_large_prime(bits)
    print(f"Generated prime p = {p}")

    # Task 1: Generate random curve coefficients A and B
    A, B = generate_random_curve(p)
    print(f"Random curve coefficients: A = {A}, B = {B}")

    # Task 2: Find a random point on the curve
    P = find_random_point(A, B, p)
    print(f"Random point P = {P}")

    # Task 3: Compute the inverse of point P
    negative_P = point_negation(P, p)
    print(f"Inverse of P: -P = {negative_P}")

    # Task 4: Add two points P and Q on the curve
    Q = find_random_point(A, B, p)
    print(f"Another random point Q = {Q}")
    R = point_addition(P, Q, A, p)
    print(f"Sum of P and Q: R = P ⊕ Q = {R}")

    # Task 5: Compute nP for a random scalar n
    n = mpz_random(random_state(random.randrange(2**32)), p)
    nP = scalar_multiplication(P, n, A, p)
    print(f"Scalar multiplication: {n} * P = {nP}")