import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Scanner;

public class RSAImplementation {

    static BigInteger p, q, n, phi, e, d;
    static BigInteger[] encryptedMessage;

    static void generateKeys() {
        SecureRandom random = new SecureRandom();

        p = BigInteger.probablePrime(512, random);
        q = BigInteger.probablePrime(512, random);

        n = p.multiply(q);
        phi = p.subtract(BigInteger.ONE)
                .multiply(q.subtract(BigInteger.ONE));

        e = BigInteger.valueOf(65537);

        while (!phi.gcd(e).equals(BigInteger.ONE)) {
            e = e.add(BigInteger.TWO);
        }

        d = e.modInverse(phi);
    }

    static void encrypt(String message) {
        byte[] bytes = message.getBytes(StandardCharsets.UTF_8);

        // Convert the complete message to a positive integer.
        BigInteger m = new BigInteger(1, bytes);

        if (m.compareTo(n) >= 0) {
            System.out.println("Message is too long for this RSA key.");
            return;
        }

        BigInteger c = m.modPow(e, n);
        encryptedMessage = new BigInteger[]{c};

        System.out.println("\n--- Encryption Successful ---");
        System.out.println("Encrypted Data (Hex): "
                + c.toString(16).toUpperCase());
    }

    static void decrypt() {
        if (encryptedMessage == null) {
            System.out.println("No encrypted message available.");
            return;
        }

        BigInteger c = encryptedMessage[0];
        BigInteger m = c.modPow(d, n);

        byte[] bytes = m.toByteArray();

        // Remove a possible leading sign byte.
        if (bytes.length > 1 && bytes[0] == 0) {
            byte[] temp = new byte[bytes.length - 1];
            System.arraycopy(bytes, 1, temp, 0, temp.length);
            bytes = temp;
        }

        String message = new String(bytes, StandardCharsets.UTF_8);

        System.out.println("\n--- Decryption Successful ---");
        System.out.println("Decrypted Message: " + message);
    }

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

        System.out.println("=== Interactive RSA Tool ===");
        System.out.println("Generating 1024-bit RSA key pair...");

        generateKeys();

        System.out.println("Done!");
        System.out.println("\nPublic Key [e]: " + e);
        System.out.println("Modulus [n]: " + n);

        while (true) {
            System.out.println("\n-----------------------------");
            System.out.println("Choose an option:");
            System.out.println("1. Encrypt a custom message");
            System.out.println("2. Decrypt the last encrypted message");
            System.out.println("3. Exit");
            System.out.print("Enter choice (1-3): ");

            int choice = sc.nextInt();
            sc.nextLine();

            switch (choice) {
                case 1:
                    System.out.print(
                            "Enter the secret message to encrypt: ");
                    String message = sc.nextLine();
                    encrypt(message);
                    break;

                case 2:
                    decrypt();
                    break;

                case 3:
                    System.out.println("Exiting RSA Tool.");
                    sc.close();
                    return;

                default:
                    System.out.println("Invalid choice!");
            }
        }
    }
}