import java.nio.charset.StandardCharsets;
import java.security.BinaryEncodable;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PEM;
import java.security.PEMDecoder;
import java.security.PEMEncoder;
import java.security.PrivateKey;
import java.security.PublicKey;
import javax.crypto.EncryptedPrivateKeyInfo;
import java.util.Arrays;
import java.util.Base64;

/**
 * JEP 538, PEM Encodings of Cryptographic Objects (third preview), against the JDK 27 API.
 *
 * <p>PEM is the "-----BEGIN ...-----" text format every key and certificate file on a Linux box is in.
 * Before this API, turning a key into PEM or back meant hand-rolled Base64 and header strings.
 * Note the interface name: JDK 26 called the common supertype {@code DEREncodable}; JDK 27 renamed it
 * {@code BinaryEncodable}. See broken/Pem26Style.java. Explained in docs/08-pem-api.md.
 * Compile and run with --enable-preview.
 */
public class PemDemo {

    public static void main(String[] args) throws Exception {
        KeyPairGenerator gen = KeyPairGenerator.getInstance("EC");
        gen.initialize(256);
        KeyPair pair = gen.generateKeyPair();

        // 1. Encode: one call, no Base64 in sight.
        String publicPem = PEMEncoder.of().encodeToString(pair.getPublic());
        String privatePem = PEMEncoder.of().encodeToString(pair.getPrivate());
        System.out.println("--- public key, encoded");
        System.out.println(firstLines(publicPem, 1) + "  ...  " + lastLine(publicPem));
        System.out.println("--- private key, encoded");
        System.out.println(firstLines(privatePem, 1) + "  ...  " + lastLine(privatePem));
        Check.that(publicPem.startsWith("-----BEGIN PUBLIC KEY-----"), "a PublicKey encodes as BEGIN PUBLIC KEY");
        Check.that(privatePem.startsWith("-----BEGIN PRIVATE KEY-----"), "a PrivateKey encodes as BEGIN PRIVATE KEY (PKCS#8)");

        // 2. Decode, asking for the type you expect.
        PublicKey publicBack = PEMDecoder.of().decode(publicPem, PublicKey.class);
        PrivateKey privateBack = PEMDecoder.of().decode(privatePem, PrivateKey.class);
        Check.that(publicBack.equals(pair.getPublic()), "the decoded public key equals the original");
        Check.that(privateBack.equals(pair.getPrivate()), "the decoded private key equals the original");

        // 3. Decode without saying what you expect, and switch on what came back.
        System.out.println("--- decode(String) returns whatever the header says it is");
        for (String pem : new String[] {publicPem, privatePem}) {
            BinaryEncodable decoded = PEMDecoder.of().decode(pem);
            String kind = switch (decoded) {
                case PublicKey k -> "PublicKey (" + k.getAlgorithm() + ")";
                case PrivateKey k -> "PrivateKey (" + k.getAlgorithm() + ")";
                default -> decoded.getClass().getName();
            };
            System.out.println(pem.lines().findFirst().orElse("") + "  ->  " + kind);
        }

        // 4. Encrypted private keys: the password goes on the encoder and the decoder.
        char[] password = "correct horse".toCharArray();
        String encryptedPem = PEMEncoder.of().withEncryption(password).encodeToString(pair.getPrivate());
        System.out.println("--- encrypted private key");
        System.out.println(firstLines(encryptedPem, 1));
        Check.that(encryptedPem.startsWith("-----BEGIN ENCRYPTED PRIVATE KEY-----"), "withEncryption produces ENCRYPTED PRIVATE KEY");

        BinaryEncodable withoutPassword = PEMDecoder.of().decode(encryptedPem);
        System.out.println("decoded with no password: " + withoutPassword.getClass().getName());
        Check.that(withoutPassword instanceof EncryptedPrivateKeyInfo, "without a password you get the still-encrypted structure back");

        PrivateKey decrypted = PEMDecoder.of().withDecryption(password).decode(encryptedPem, PrivateKey.class);
        Check.that(decrypted.equals(pair.getPrivate()), "with the password you get the original private key");

        // 5. A PEM type the JDK does not know: it round-trips as a raw PEM object.
        //    The trap: PEM's content is the Base64 TEXT, not the payload. Hand it raw bytes and they are written verbatim.
        byte[] payload = "hello, pem".getBytes(StandardCharsets.UTF_8);
        String wrong = PEMEncoder.of().encodeToString(new PEM("ANKURM DEMO", payload));
        System.out.println("--- PEM(type, byte[]) does NOT Base64-encode: the bytes go in as they are");
        System.out.println(wrong.strip());
        Check.that(wrong.contains("hello, pem"), "raw bytes appear verbatim between the header and footer");

        String base64 = Base64.getEncoder().encodeToString(payload);
        String right = PEMEncoder.of().encodeToString(new PEM("ANKURM DEMO", base64));
        System.out.println("--- give it the Base64 text and it round-trips");
        System.out.println(right.strip());
        BinaryEncodable raw = PEMDecoder.of().decode(right);
        Check.that(raw instanceof PEM p && p.type().equals("ANKURM DEMO")
                        && new String(p.content(), StandardCharsets.US_ASCII).strip().equals(base64)
                        && Arrays.equals(p.decode(), payload),
                "type, Base64 text (content()) and decoded payload (decode()) all survive");
    }

    static String firstLines(String s, int n) {
        return s.lines().limit(n).reduce((a, b) -> a + "\n" + b).orElse("");
    }

    static String lastLine(String s) {
        return s.lines().reduce((a, b) -> b).orElse("");
    }
}