import java.math.BigInteger;
import java.util.Scanner;

public class rabin{
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        System.out.print("Enter prime p (p mod 4 = 3): ");
        BigInteger p = sc.nextBigInteger();

        System.out.print("Enter prime q (q mod 4 = 3): ");
        BigInteger q = sc.nextBigInteger();

        if (!(p.isProbablePrime(100)
                && q.isProbablePrime(100)
                && p.mod(BigInteger.valueOf(4)).equals(BigInteger.valueOf(3))
                && q.mod(BigInteger.valueOf(4)).equals(BigInteger.valueOf(3)))) {
            System.out.println("Invalid primes.");
            sc.close();
            return;
        }

        BigInteger n = p.multiply(q);
        System.out.println("Public key (n = p*q) = " + n);

        System.out.print("Enter plaintext message m (m < n): ");
        BigInteger m = sc.nextBigInteger();

        BigInteger c = m.modPow(BigInteger.valueOf(2), n);
        System.out.println("Encrypted Ciphertext C = " + c);

        BigInteger mp =
                c.modPow(p.add(BigInteger.ONE).divide(BigInteger.valueOf(4)), p);
        BigInteger mq =
                c.modPow(q.add(BigInteger.ONE).divide(BigInteger.valueOf(4)), q);

        BigInteger[] egcd = extendedGCD(p, q);
        BigInteger yp = egcd[1];
        BigInteger yq = egcd[2];

        BigInteger r1 = yp.multiply(p).multiply(mq)
                .add(yq.multiply(q).multiply(mp)).mod(n);
        BigInteger r2 = n.subtract(r1).mod(n);

        BigInteger r3 = yp.multiply(p).multiply(mq.negate())
                .add(yq.multiply(q).multiply(mp)).mod(n);
        BigInteger r4 = n.subtract(r3).mod(n);

        System.out.println("\nPossible plaintext roots:");
        System.out.println("M1 = " + r1);
        System.out.println("M2 = " + r2);
        System.out.println("M3 = " + r3);
        System.out.println("M4 = " + r4);

        if (m.equals(r1) || m.equals(r2) || m.equals(r3) || m.equals(r4)) {
            System.out.println("\nOriginal plaintext recovered: " + m);
        } else {
            System.out.println("\nOriginal plaintext not found.");
        }

        sc.close();
    }

    static BigInteger[] extendedGCD(BigInteger a, BigInteger b) {
        if (b.equals(BigInteger.ZERO)) {
            return new BigInteger[]{
                    a, BigInteger.ONE, BigInteger.ZERO
            };
        }
        BigInteger[] result = extendedGCD(b, a.mod(b));
        BigInteger gcd = result[0];
        BigInteger x = result[2];
        BigInteger y = result[1].subtract(a.divide(b).multiply(result[2]));

        return new BigInteger[]{gcd, x, y};
    }
}