package com.zb.mom.ww.mxgateway.client;

import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import io.grpc.ManagedChannel;
import io.grpc.Server;
import io.grpc.StatusRuntimeException;
import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
import io.grpc.stub.StreamObserver;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.file.Files;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.time.Duration;
import java.util.Base64;
import java.util.concurrent.TimeUnit;
import javax.net.ssl.SSLException;
import mxaccess_gateway.v1.MxAccessGatewayGrpc;
import mxaccess_gateway.v1.MxaccessGateway.OpenSessionReply;
import mxaccess_gateway.v1.MxaccessGateway.OpenSessionRequest;
import mxaccess_gateway.v1.MxaccessGateway.ProtocolStatus;
import mxaccess_gateway.v1.MxaccessGateway.ProtocolStatusCode;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

/**
 * Verifies that the Java client connects to a Netty TLS server with a
 * self-signed certificate when no CA is pinned (lenient default), and that
 * setting {@code requireCertificateValidation(true)} causes a TLS failure.
 *
 * <p>A self-signed certificate is generated using {@code keytool} (always
 * available in the JDK) to avoid dependencies on internal JDK APIs or
 * BouncyCastle, and so the test works on all JDK versions used by the project.
 */
final class MxGatewayClientTlsTests {

    private Server server;
    private int port;
    private File certPemFile;
    private File keyPemFile;
    private File keystoreFile;

    @BeforeEach
    void startTlsServer() throws Exception {
        keystoreFile = File.createTempFile("gw-test-ks", ".p12");
        certPemFile = File.createTempFile("gw-test-cert", ".pem");
        keyPemFile = File.createTempFile("gw-test-key", ".pem");

        // keytool refuses to write to a pre-existing (even empty) file; delete it first.
        keystoreFile.delete();

        // Use keytool to generate a self-signed PKCS12 keystore.
        String keytool = ProcessHandle.current().info().command()
                .map(cmd -> cmd.replace("java", "keytool"))
                .orElse("keytool");
        // Fall back to just "keytool" on PATH if the resolved path doesn't exist.
        if (!new File(keytool).exists()) {
            keytool = "keytool";
        }
        Process p = new ProcessBuilder(
                keytool,
                "-genkeypair",
                "-alias", "server",
                "-keyalg", "RSA",
                "-keysize", "2048",
                "-sigalg", "SHA256withRSA",
                "-validity", "1",
                "-dname", "CN=localhost",
                "-storetype", "PKCS12",
                "-storepass", "changeit",
                "-keypass", "changeit",
                "-keystore", keystoreFile.getAbsolutePath())
                .redirectErrorStream(true)
                .start();
        int exit = p.waitFor();
        if (exit != 0) {
            String out = new String(p.getInputStream().readAllBytes());
            throw new IllegalStateException("keytool failed (exit " + exit + "): " + out);
        }

        // Export cert and private key from the PKCS12 keystore to PEM files.
        KeyStore ks = KeyStore.getInstance("PKCS12");
        try (var is = Files.newInputStream(keystoreFile.toPath())) {
            ks.load(is, "changeit".toCharArray());
        }
        X509Certificate cert = (X509Certificate) ks.getCertificate("server");
        PrivateKey privateKey = (PrivateKey) ks.getKey("server", "changeit".toCharArray());

        try (FileOutputStream out = new FileOutputStream(certPemFile)) {
            out.write("-----BEGIN CERTIFICATE-----\n".getBytes());
            out.write(Base64.getMimeEncoder(64, new byte[]{'\n'}).encode(cert.getEncoded()));
            out.write("\n-----END CERTIFICATE-----\n".getBytes());
        }
        try (FileOutputStream out = new FileOutputStream(keyPemFile)) {
            out.write("-----BEGIN PRIVATE KEY-----\n".getBytes());
            out.write(Base64.getMimeEncoder(64, new byte[]{'\n'}).encode(privateKey.getEncoded()));
            out.write("\n-----END PRIVATE KEY-----\n".getBytes());
        }

        server = NettyServerBuilder
                .forAddress(new InetSocketAddress("127.0.0.1", 0))
                .sslContext(GrpcSslContexts.forServer(certPemFile, keyPemFile).build())
                .addService(new MinimalGatewayService())
                .build()
                .start();
        port = server.getPort();
    }

    @AfterEach
    void stopTlsServer() throws InterruptedException {
        if (server != null) {
            server.shutdown();
            server.awaitTermination(5, TimeUnit.SECONDS);
        }
        if (certPemFile != null) {
            certPemFile.delete();
        }
        if (keyPemFile != null) {
            keyPemFile.delete();
        }
        if (keystoreFile != null) {
            keystoreFile.delete();
        }
    }

    @Test
    void connectsToSelfSignedServer_WhenRequireCertificateValidationIsFalse() throws SSLException {
        // Default options — requireCertificateValidation defaults to false.
        MxGatewayClientOptions options = MxGatewayClientOptions.builder()
                .endpoint("127.0.0.1:" + port)
                .apiKey("test-key")
                .connectTimeout(Duration.ofSeconds(5))
                .callTimeout(Duration.ofSeconds(5))
                .build();

        ManagedChannel channel = MxGatewayClient.createChannelForTests(options);
        try {
            MxAccessGatewayGrpc.MxAccessGatewayBlockingStub stub =
                    MxAccessGatewayGrpc.newBlockingStub(channel);
            OpenSessionReply reply = stub.openSession(
                    OpenSessionRequest.newBuilder()
                            .setClientSessionName("tls-test")
                            .build());
            assertTrue(reply.getProtocolStatus().getCode()
                    == ProtocolStatusCode.PROTOCOL_STATUS_CODE_OK);
        } finally {
            channel.shutdownNow();
        }
    }

    @Test
    void failsToConnect_WhenRequireCertificateValidationIsTrue() throws SSLException {
        MxGatewayClientOptions options = MxGatewayClientOptions.builder()
                .endpoint("127.0.0.1:" + port)
                .apiKey("test-key")
                .requireCertificateValidation(true)
                .connectTimeout(Duration.ofSeconds(5))
                .callTimeout(Duration.ofSeconds(5))
                .build();

        ManagedChannel channel = MxGatewayClient.createChannelForTests(options);
        try {
            MxAccessGatewayGrpc.MxAccessGatewayBlockingStub stub =
                    MxAccessGatewayGrpc.newBlockingStub(channel);
            assertThrows(StatusRuntimeException.class, () ->
                    stub.openSession(OpenSessionRequest.newBuilder()
                            .setClientSessionName("tls-strict-test")
                            .build()));
        } finally {
            channel.shutdownNow();
        }
    }

    /** Minimal gateway stub that succeeds any OpenSession call. */
    private static final class MinimalGatewayService
            extends MxAccessGatewayGrpc.MxAccessGatewayImplBase {
        @Override
        public void openSession(
                OpenSessionRequest request,
                StreamObserver<OpenSessionReply> responseObserver) {
            responseObserver.onNext(OpenSessionReply.newBuilder()
                    .setSessionId("tls-test-session")
                    .setProtocolStatus(ProtocolStatus.newBuilder()
                            .setCode(ProtocolStatusCode.PROTOCOL_STATUS_CODE_OK)
                            .build())
                    .build());
            responseObserver.onCompleted();
        }
    }
}
