nomyo4J/src/main/java/ai/nomyo/SecureCompletionClient.java

package ai.nomyo;

import ai.nomyo.errors.*;
import ai.nomyo.util.PEMConverter;
import ai.nomyo.util.Pass2Key;
import lombok.Getter;
import lombok.Setter;

import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.math.BigInteger;
import java.net.*;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.PosixFilePermission;
import java.nio.file.attribute.PosixFilePermissions;
import java.security.*;
import java.security.spec.*;
import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutionException;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;

import java.util.concurrent.locks.ReentrantLock;

/**
 * Low-level client: key management, hybrid encryption, HTTP with retry, response decryption. Used by {@link SecureChatCompletion}.
 */
public class SecureCompletionClient {

    /**
     * NOMYO router base URL (trailing slash stripped).
     */
    @Getter
    private final String routerUrl;

    /**
     * Permit HTTP (non-HTTPS) URLs.
     */
    @Getter
    private final boolean allowHttp;

    /**
     * RSA key size in bits ({@link Constants#RSA_KEY_SIZE}).
     */
    @Getter
    private final int keySize;

    /**
     * Max retries on retryable errors.
     */
    @Getter
    private final int maxRetries;

    /**
     * Secure memory operations active.
     */
    @Getter
    private final boolean useSecureMemory;

    /**
     * Lock for double-checked key initialization.
     */
    private final ReentrantLock keyInitLock = new ReentrantLock();
    private final HttpClient httpClient;
    /**
     * RSA private key ({@code null} until loaded/generated).
     */
    @Getter
    private PrivateKey privateKey;
    /**
     * PEM-encoded public key ({@code null} until loaded/generated).
     */
    @Getter
    private String publicPemKey;
    /**
     * Keys initialized.
     */
    private volatile boolean keysInitialized = false;

    /**
     * Default settings: {@code https://api.nomyo.ai}, HTTPS-only, secure memory, 2 retries.
     */
    public SecureCompletionClient() {
        this(Constants.DEFAULT_BASE_URL, false, true, Constants.DEFAULT_MAX_RETRIES);
    }

    /**
     * @param routerUrl    NOMYO router base URL
     * @param allowHttp    permit HTTP URLs
     * @param secureMemory enable memory locking/zeroing
     * @param maxRetries   retries on retryable errors
     */
    public SecureCompletionClient(String routerUrl, boolean allowHttp, boolean secureMemory, int maxRetries) {
        this.routerUrl = routerUrl != null ? routerUrl.replaceAll("/+$", "") : Constants.DEFAULT_BASE_URL;
        this.allowHttp = allowHttp;
        this.useSecureMemory = secureMemory;
        this.keySize = Constants.RSA_KEY_SIZE;
        this.maxRetries = maxRetries;
        this.httpClient = HttpClient.newHttpClient();
    }

    private static String readFileContent(String filePath) throws IOException {
        return Files.readString(Path.of(filePath));
    }

    /**
     * Generates a 4096-bit RSA key pair (exponent 65537). Saves to disk if {@code saveToFile}.
     */
    public void generateKeys(boolean saveToFile, String keyDir, String password) throws SecurityError {
        try {
            KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");
            generator.initialize(new RSAKeyGenParameterSpec(Constants.RSA_KEY_SIZE, BigInteger.valueOf(Constants.RSA_PUBLIC_EXPONENT)));

            KeyPair pair = generator.generateKeyPair();

            String privatePem = PEMConverter.toPEM(pair.getPrivate().getEncoded(), true);
            String publicPem = PEMConverter.toPEM(pair.getPublic().getEncoded(), false);

            if (saveToFile) {
                Path keyFolder = Path.of(keyDir);
                if (!Files.exists(keyFolder)) {
                    try {
                        Files.createDirectories(keyFolder);
                    } catch (IOException e) {
                        throw new IOException("Failed to create key directory: " + keyDir, e);
                    }
                }

                Path privateKeyPath = Path.of(keyDir, Constants.DEFAULT_PRIVATE_KEY_FILE);
                if (!Files.exists(privateKeyPath)) {
                    Set<PosixFilePermission> filePermissions = PosixFilePermissions.fromString(Constants.PRIVATE_KEY_FILE_MODE);
                    Files.createFile(privateKeyPath, PosixFilePermissions.asFileAttribute(filePermissions));

                    try (var writer = Files.newBufferedWriter(privateKeyPath)) {
                        if (password == null || password.isEmpty()) {
                            System.out.println("WARNING: Saving keys in plaintext!");
                        } else {
                            try {
                                privatePem = Pass2Key.encrypt("AES/GCM/NoPadding", privatePem, password);
                            } catch (NoSuchPaddingException | IllegalBlockSizeException | BadPaddingException |
                                     InvalidKeyException | SecurityError e) {
                                throw new SecurityError("Failed to encrypt private key with password: " + e.getMessage(), e);
                            }
                        }
                        writer.write(privatePem);
                    }
                }

                Path publicKeyPath = Path.of(keyDir, Constants.DEFAULT_PUBLIC_KEY_FILE);
                if (!Files.exists(publicKeyPath)) {
                    Set<PosixFilePermission> publicPermissions = PosixFilePermissions.fromString(Constants.PUBLIC_KEY_FILE_MODE);
                    Files.createFile(publicKeyPath, PosixFilePermissions.asFileAttribute(publicPermissions));

                    try (var writer = Files.newBufferedWriter(publicKeyPath)) {
                        writer.write(publicPem);
                    }
                }
            }

            this.privateKey = pair.getPrivate();
            this.publicPemKey = publicPem;

        } catch (SecurityError e) {
            throw e;
        } catch (NoSuchAlgorithmException e) {
            throw new SecurityError("RSA algorithm not available: " + e.getMessage(), e);
        } catch (InvalidAlgorithmParameterException e) {
            throw new SecurityError("Invalid RSA key generation parameters: " + e.getMessage(), e);
        } catch (IOException e) {
            throw new RuntimeException("Failed to save keys: " + e.getMessage(), e);
        }
    }

    /**
     * Generates a 4096-bit RSA key pair and saves to the default directory.
     */
    public void generateKeys(boolean saveToFile) throws SecurityError {
        generateKeys(saveToFile, Constants.DEFAULT_KEY_DIR, null);
    }

    /**
     * Loads RSA private key from disk. If {@code publicPemKeyPath} is {@code null}, derives public key.
     * Validates key size >= {@link Constants#MIN_RSA_KEY_SIZE} bits.
     *
     * @param privateKeyPath   private key PEM path
     * @param publicPemKeyPath optional public key PEM path
     * @param password         optional password for encrypted private key
     * @throws SecurityError if key file not found, unreadable, or decryption fails
     */
    public void loadKeys(String privateKeyPath, String publicPemKeyPath, String password) throws SecurityError {
        Path keyPath = Path.of(privateKeyPath);
        if (!Files.exists(keyPath)) {
            throw new SecurityError("Private key file not found: " + privateKeyPath);
        }

        String keyContent;
        if (password != null && !password.isEmpty()) {
            try {
                keyContent = readFileContent(privateKeyPath);
            } catch (IOException e) {
                throw new SecurityError("Failed to read private key file: " + e.getMessage(), e);
            }

            try {
                keyContent = Pass2Key.decrypt("AES/GCM/NoPadding", keyContent, password);
            } catch (NoSuchPaddingException | NoSuchAlgorithmException | BadPaddingException |
                     IllegalBlockSizeException | InvalidAlgorithmParameterException | InvalidKeyException | SecurityError e) {
                throw new SecurityError("Failed to decrypt private key with provided password: " + e.getMessage(), e);
            }
        } else {
            try {
                keyContent = readFileContent(privateKeyPath);
            } catch (IOException e) {
                throw new SecurityError("Failed to read private key file: " + e.getMessage(), e);
            }
        }

        try {
            this.privateKey = Pass2Key.convertStringToPrivateKey(keyContent);
        } catch (Exception e) {
            throw new SecurityError("Failed to load private key: " + e.getMessage(), e);
        }
    }

    /**
     * Loads RSA private key from disk, deriving public key.
     *
     * @throws SecurityError if key file not found, unreadable, or decryption fails
     */
    public void loadKeys(String privateKeyPath, String password) throws SecurityError {
        loadKeys(privateKeyPath, null, password);
    }

    /**
     * GET {@code {routerUrl}/pki/public_key}. Returns server PEM public key.
     * Errors are propagated as {@link CompletionException} wrappers around the underlying checked exception.
     */
    public CompletableFuture<String> fetchServerPublicKey() {
        if (!this.routerUrl.startsWith("https://")) {
            if (!this.allowHttp) {
                return CompletableFuture.failedFuture(new SecurityError("Server public key must be fetched over HTTPS to prevent MITM attacks."));
            } else {
                System.out.println("Fetching server public key over HTTP (local dev mode)");
            }
        }

        URI url;

        try {
            url = new URI(this.routerUrl + Constants.PKI_PUBLIC_KEY_PATH);
        } catch (URISyntaxException e) {
            return CompletableFuture.failedFuture(new CompletionException(new APIConnectionError("Invalid URI: " + e.getMessage(), e)));
        }

        HttpRequest request = HttpRequest.newBuilder(url).timeout(Duration.ofSeconds(Constants.DEFAULT_TIMEOUT_SECONDS)).GET().build();

        return this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()).thenApply(response -> {
            if (response.statusCode() != 200) {
                throw new CompletionException(new APIConnectionError("Could not fetch server public key!"));
            }
            return response.body();
        }).thenApply(body -> {
            if (!PEMConverter.validatePEM(body)) {
                throw new CompletionException(new SecurityError("Server returned invalid PEM key format (possible MITM attack)"));
            }
            return body;
        });
    }

    /**
     * Hybrid encryption: AES-256-GCM for payload, RSA-OAEP-SHA256 for AES key wrapping.
     *
     * @param payload OpenAI-compatible chat parameters
     * @return encrypted bytes (JSON package)
     * @throws SecurityError if encryption fails or keys not loaded
     */
    public CompletableFuture<byte[]> encryptPayload(Map<String, Object> payload) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                ensureKeys(null);
            } catch (SecurityError e) {
                throw new CompletionException(new SecurityError("Failed to ensure keys are initialized: " + e.getMessage(), e));
            }

            if (this.privateKey == null) {
                throw new CompletionException(new SecurityError("Private key not available for encryption"));
            }

            // Generate AES key
            KeyGenerator keyGen;

            try {
                keyGen = KeyGenerator.getInstance("AES");
                keyGen.init(Constants.AES_KEY_SIZE * 8);
            } catch (NoSuchAlgorithmException e) {
                throw new CompletionException(new SecurityError("AES key generation not available: " + e.getMessage(), e));
            }

            Key aesKey = keyGen.generateKey();

            // Serialize payload to JSON
            Gson gson = new Gson();
            String payloadJson = gson.toJson(payload);
            byte[] payloadBytes = payloadJson.getBytes(StandardCharsets.UTF_8);

            // Encrypt
            return doEncrypt(payloadBytes, aesKey).join();
        });
    }

    /**
     * Core hybrid encryption: AES-256-GCM encrypts {@code payloadBytes} with {@code aesKey}.
     */
    public CompletableFuture<byte[]> doEncrypt(byte[] payloadBytes, Key aesKey) {
        return CompletableFuture.supplyAsync(() -> {
            SecureRandom random = new SecureRandom();
            byte[] nonce = new byte[Constants.GCM_NONCE_SIZE];
            random.nextBytes(nonce);

            Cipher cipher = null;
            try {
                cipher = Cipher.getInstance("AES/GCM/NoPadding");
                cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(aesKey.getEncoded(), "AES"), new GCMParameterSpec(Constants.GCM_TAG_SIZE * Byte.SIZE, nonce));
            } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidAlgorithmParameterException |
                     InvalidKeyException e) {
                throw new RuntimeException(new SecurityError("AES-GCM cipher initialization failed: " + e.getMessage(), e));
            }

            byte[] ciphertext;

            try {
                ciphertext = cipher.doFinal(payloadBytes);
            } catch (IllegalBlockSizeException | BadPaddingException e) {
                throw new RuntimeException(new SecurityError("AES-GCM encryption failed: " + e.getMessage(), e));
            }

            String serverPEM;

            try {
                serverPEM = fetchServerPublicKey().get();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new RuntimeException(new SecurityError("Encryption interrupted while fetching server public key", e));
            } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof SecurityError) {
                    throw new RuntimeException((SecurityError) cause);
                }
                throw new RuntimeException(new SecurityError("Failed to fetch server public key: " + cause.getMessage(), cause));
            }

            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(PEMConverter.fromPEM(serverPEM));

            PublicKey serverPublicKey;

            try {
                serverPublicKey = KeyFactory.getInstance("RSA").generatePublic(keySpec);
            } catch (InvalidKeySpecException | NoSuchAlgorithmException e) {
                throw new RuntimeException(new SecurityError("RSA key factory failed to parse server public key: " + e.getMessage(), e));
            }

            Cipher rsa;
            try {
                rsa = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
                rsa.init(Cipher.ENCRYPT_MODE, serverPublicKey);
            } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException e) {
                throw new RuntimeException(new SecurityError("RSA-OAEP cipher initialization failed: " + e.getMessage(), e));
            }
            byte[] encryptedAESKey;
            try {
                encryptedAESKey = rsa.doFinal(aesKey.getEncoded());
            } catch (IllegalBlockSizeException | BadPaddingException e) {
                throw new RuntimeException(new SecurityError("RSA-OAEP key wrapping failed: " + e.getMessage(), e));
            }

            byte[] tag = Arrays.copyOfRange(ciphertext, ciphertext.length - Constants.GCM_TAG_SIZE, ciphertext.length);

            EncryptedRequest request = new EncryptedRequest();

            request.setVersion(Constants.PROTOCOL_VERSION);
            request.setAlgorithm(Constants.HYBRID_ALGORITHM);
            request.setEncryptedPayload(new EncryptedRequest.EncryptedPayload(Base64.getEncoder().encodeToString(ciphertext), Base64.getEncoder().encodeToString(nonce), Base64.getEncoder().encodeToString(tag)));
            request.setEncryptedAESKey(Base64.getEncoder().encodeToString(encryptedAESKey));
            request.setKeyAlgorithm(Constants.KEY_WRAP_ALGORITHM);
            request.setPayloadAlgorithm(Constants.PAYLOAD_ALGORITHM);

            return request.toJson().getBytes(StandardCharsets.UTF_8);
        });
    }

    /**
     * encrypt → POST {routerUrl}/v1/chat/secure_completion → retry → decrypt → return.
     * <p>Headers: Content-Type=octet-stream, X-Payload-ID, X-Public-Key, Authorization (Bearer), X-Security-Tier.
     * Retryable: 429, 500, 502, 503, 504 + network errors. Backoff: 2^(attempt-1)s.
     * <p>Status mapping: 200→return, 400→InvalidRequestError, 401→AuthenticationError, 403→ForbiddenError,
     * 404→APIError, 429→RateLimitError, 500→ServerError, 503→ServiceUnavailableError,
     * 502/504→APIError(retryable), network→APIConnectionError.
     *
     * @param payload      OpenAI-compatible chat parameters
     * @param payloadId    unique payload identifier
     * @param apiKey       optional API key
     * @param securityTier optional: "standard", "high", "maximum"
     * @return decrypted response map
     * @throws SecurityError           encryption/decryption failure
     * @throws APIConnectionError      network error
     * @throws InvalidRequestError     HTTP 400
     * @throws AuthenticationError     HTTP 401
     * @throws ForbiddenError          HTTP 403
     * @throws RateLimitError          HTTP 429
     * @throws ServerError             HTTP 500
     * @throws ServiceUnavailableError HTTP 503
     * @throws APIError                other errors
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey, String securityTier) {
        return CompletableFuture.supplyAsync(() -> {
            // Validate security tier if provided
            if (securityTier != null && !Constants.VALID_SECURITY_TIERS.contains(securityTier)) {
                throw new CompletionException(new ValueError(
                    "Invalid security_tier: '" + securityTier + "'. Must be one of: " + Constants.VALID_SECURITY_TIERS));
            }

            try {
                ensureKeys(null);
            } catch (SecurityError e) {
                throw new CompletionException(new SecurityError("Failed to ensure keys: " + e.getMessage(), e));
            }

            // Step 1: Encrypt payload
            byte[] encryptedPayload;
            try {
                encryptedPayload = encryptPayload(payload).get();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new CompletionException(new APIConnectionError("Encryption interrupted"));
            } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof SecurityError) {
                    throw new CompletionException((SecurityError) cause);
                }
                throw new CompletionException(new SecurityError("Encryption failed: " + cause.getMessage(), cause));
            }

            // Step 2: Prepare headers
            URI url;
            try {
                url = new URI(this.routerUrl + Constants.SECURE_COMPLETION_PATH);
            } catch (URISyntaxException e) {
                throw new CompletionException(new APIConnectionError("Invalid URL: " + e.getMessage(), e));
            }

            HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(url)
                .timeout(Duration.ofSeconds(Constants.DEFAULT_TIMEOUT_SECONDS))
                .header("Content-Type", Constants.CONTENT_TYPE_OCTET_STREAM)
                .header(Constants.HEADER_PAYLOAD_ID, payloadId)
                .header(Constants.HEADER_PUBLIC_KEY, urlEncodePublicKey(this.publicPemKey))
                .POST(HttpRequest.BodyPublishers.ofByteArray(encryptedPayload));

            if (apiKey != null && !apiKey.isEmpty()) {
                requestBuilder.header("Authorization", Constants.AUTHORIZATION_BEARER_PREFIX + apiKey);
            }
            if (securityTier != null) {
                requestBuilder.header(Constants.HEADER_SECURITY_TIER, securityTier);
            }

            HttpRequest request = requestBuilder.build();

            // Step 3: Send request with retry
            Exception lastExc = new APIConnectionError("Request failed");
            int retryableCodes = 0;
            for (int attempt = 0; attempt <= this.maxRetries; attempt++) {
                if (attempt > 0) {
                    long delay = (long) Math.pow(2, attempt - 1);
                    try {
                        Thread.sleep(delay * 1000);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                        throw new CompletionException(new APIConnectionError("Retry interrupted"));
                    }
                }

                try {
                    HttpResponse<byte[]> response = this.httpClient.send(request, HttpResponse.BodyHandlers.ofByteArray());
                    int statusCode = response.statusCode();

                    if (statusCode == 200) {
                        // Step 4: Decrypt response
                        return decryptResponse(response.body(), payloadId).get();
                    } else if (statusCode == 400) {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> errorDetails = Map.of();
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) errorDetails = parsed;
                        } catch (Exception ignored) {
                        }
                        String detail = errorDetails.containsKey("detail") ? errorDetails.get("detail").toString() : "Unknown error";
                        throw new CompletionException(new InvalidRequestError("Bad request: " + detail, 400, errorDetails));
                    } else if (statusCode == 401) {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> errorDetails = Map.of();
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) errorDetails = parsed;
                        } catch (Exception ignored) {
                        }
                        String detail = errorDetails.containsKey("detail") ? errorDetails.get("detail").toString() : "Invalid API key or authentication failed";
                        throw new CompletionException(new AuthenticationError(detail, 401, errorDetails));
                    } else if (statusCode == 403) {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> errorDetails = Map.of();
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) errorDetails = parsed;
                        } catch (Exception ignored) {
                        }
                        String detail = errorDetails.containsKey("detail") ? errorDetails.get("detail").toString() : "Model not allowed for the requested security tier";
                        throw new CompletionException(new ForbiddenError("Forbidden: " + detail, 403, errorDetails));
                    } else if (statusCode == 404) {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> errorDetails = Map.of();
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) errorDetails = parsed;
                        } catch (Exception ignored) {
                        }
                        String detail = errorDetails.containsKey("detail") ? errorDetails.get("detail").toString() : "Secure inference not enabled";
                        throw new CompletionException(new APIError("Endpoint not found: " + detail, 404, errorDetails));
                    } else if (Constants.RETRYABLE_STATUS_CODES.contains(statusCode)) {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> error = Map.of();
                        String detailMsg = "unknown";
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) {
                                error = parsed;
                                if (error.containsKey("detail")) detailMsg = error.get("detail").toString();
                            }
                        } catch (Exception ignored) {
                        }

                        if (statusCode == 429) {
                            lastExc = new RateLimitError("Rate limit exceeded: " + detailMsg, 429, error);
                        } else if (statusCode == 500) {
                            lastExc = new ServerError("Server error: " + detailMsg, 500, error);
                        } else if (statusCode == 503) {
                            lastExc = new ServiceUnavailableError("Service unavailable: " + detailMsg, 503, error);
                        } else {
                            lastExc = new APIError("Unexpected status code: " + statusCode + " " + detailMsg, statusCode, error);
                        }

                        if (attempt < this.maxRetries) {
                            continue;
                        }
                        throw new CompletionException(lastExc);
                    } else {
                        String body = new String(response.body(), StandardCharsets.UTF_8);
                        Map<String, Object> errorDetails = Map.of();
                        String detailMsg = "unknown";
                        try {
                            @SuppressWarnings("unchecked")
                            Map<String, Object> parsed = (Map<String, Object>) new Gson().fromJson(body, Object.class);
                            if (parsed != null) {
                                errorDetails = parsed;
                                if (errorDetails.containsKey("detail")) detailMsg = errorDetails.get("detail").toString();
                            }
                        } catch (Exception ignored) {
                        }
                        throw new CompletionException(new APIError("Unexpected status code: " + statusCode + " " + detailMsg, statusCode, errorDetails));
                    }
                } catch (CompletionException e) {
                    Throwable cause = e.getCause();
                    if (cause instanceof InvalidRequestError || cause instanceof AuthenticationError ||
                        cause instanceof ForbiddenError || cause instanceof SecurityError) {
                        throw e;
                    }
                    if (cause instanceof RateLimitError || cause instanceof ServerError || cause instanceof ServiceUnavailableError) {
                        lastExc = (Exception) cause;
                        if (attempt < this.maxRetries) {
                            continue;
                        }
                        throw new CompletionException(lastExc);
                    }
                    if (cause instanceof APIError apiError) {
                        if (Constants.RETRYABLE_STATUS_CODES.contains(apiError.getStatusCode())) {
                            lastExc = apiError;
                            if (attempt < this.maxRetries) {
                                continue;
                            }
                            throw new CompletionException(lastExc);
                        }
                        throw e;
                    }
                    lastExc = new APIConnectionError("Failed to connect to router: " + e.getMessage());
                    if (attempt < this.maxRetries) {
                        continue;
                    }
                    throw new CompletionException(lastExc);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new CompletionException(new APIConnectionError("Request interrupted"));
                } catch (ExecutionException e) {
                    throw new CompletionException(new APIConnectionError("Request failed: " + e.getMessage(), e));
                } catch (IOException e) {
                    lastExc = new APIConnectionError("IO error: " + e.getMessage(), e);
                    if (attempt < this.maxRetries) {
                        continue;
                    }
                    throw new CompletionException(lastExc);
                }
            }

            throw new CompletionException(lastExc);
        });
    }

    /**
     * Without security tier.
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId, String apiKey) {
        return sendSecureRequest(payload, payloadId, apiKey, null);
    }

    /**
     * No API key or security tier.
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(Map<String, Object> payload, String payloadId) {
        return sendSecureRequest(payload, payloadId, null, null);
    }

    /**
     * Thread-safe key init via double-checked locking. Loads from disk if {@code keyDir} set, else generates.
     *
     * @param keyDir key directory or {@code null} for ephemeral
     */
    public void ensureKeys(String keyDir) throws SecurityError {
        if (keysInitialized) return;
        keyInitLock.lock();
        try {
            if (keysInitialized) return;
            generateKeys(keyDir != null && !keyDir.isEmpty());
            keysInitialized = true;
        } finally {
            keyInitLock.unlock();
        }
    }

    /**
     * Validates RSA key size >= {@link Constants#MIN_RSA_KEY_SIZE} bits.
     */
    public void validateRsaKey(PrivateKey key) throws SecurityError {
        if (key == null) {
            throw new SecurityError("RSA key is null");
        }
        int keySize = extractKeySize(key);

        if (keySize < Constants.MIN_RSA_KEY_SIZE) {
            throw new SecurityError("RSA key size " + keySize + " bits is below minimum " + Constants.MIN_RSA_KEY_SIZE + " bits");
        }
    }

    private int extractKeySize(PrivateKey key) {
        try {
            var kf = KeyFactory.getInstance("RSA");
            try {
                var crtSpec = kf.getKeySpec(key, RSAPrivateCrtKeySpec.class);
                return crtSpec.getModulus().bitLength();
            } catch (Exception ignored) {
                var privSpec = kf.getKeySpec(key, RSAPrivateKeySpec.class);
                return privSpec.getModulus().bitLength();
            }
        } catch (Exception ignored) {
            if (key.getEncoded() != null) {
                return key.getEncoded().length * 8;
            }
        }
        return 0;
    }

    /**
     * Maps HTTP status code to exception (200→null).
     */
    public Exception mapHttpStatus(int statusCode, String responseBody) {
        String message = responseBody != null ? responseBody : "no body";
        Map<String, Object> errorDetails = responseBody != null ? Map.of("response", responseBody) : Map.of();
        return switch (statusCode) {
            case 200 -> null;
            case 400 -> new InvalidRequestError("Invalid request: " + message, statusCode, errorDetails);
            case 401 -> new AuthenticationError("Authentication failed: " + message, statusCode, errorDetails);
            case 403 -> new ForbiddenError("Access forbidden: " + message, statusCode, errorDetails);
            case 404 -> new APIError("Not found: " + message, statusCode, errorDetails);
            case 429 -> new RateLimitError("Rate limit exceeded: " + message, statusCode, errorDetails);
            case 500 -> new ServerError("Internal server error: " + message, statusCode, errorDetails);
            case 503 -> new ServiceUnavailableError("Service unavailable: " + message, statusCode, errorDetails);
            case 502, 504 -> new APIError("Gateway error: " + message, statusCode, errorDetails);
            default -> new APIError("Unexpected status " + statusCode + ": " + message, statusCode, errorDetails);
        };
    }

    /**
     * URL-encodes PEM key for {@code X-Public-Key} header.
     */
    public String urlEncodePublicKey(String pemKey) {
        return URLEncoder.encode(pemKey, StandardCharsets.UTF_8);
    }

    /**
     * Closes the HTTP client and clears keys from memory.
     */
    public void close() {
        this.httpClient.close();
        this.privateKey = null;
        this.publicPemKey = null;
        this.keysInitialized = false;
    }

    /**
     * Decrypts server response.
     */
    public CompletableFuture<Map<String, Object>> decryptResponse(byte[] encryptedResponse, String payloadId) {
        return CompletableFuture.supplyAsync(() -> {
            if (encryptedResponse == null || encryptedResponse.length == 0) {
                throw new CompletionException(new ValueError("Empty encrypted response"));
            }

            String jsonResponse;
            try {
                jsonResponse = new String(encryptedResponse, StandardCharsets.UTF_8);
            } catch (Exception e) {
                throw new CompletionException(new ValueError("Invalid encrypted package format: not valid UTF-8"));
            }

            Gson gson = new Gson();
            JsonObject packageJson;
            try {
                packageJson = JsonParser.parseString(jsonResponse).getAsJsonObject();
            } catch (Exception e) {
                throw new CompletionException(new ValueError("Invalid encrypted package format: malformed JSON"));
            }

            // Validate required fields
            String[] requiredFields = {"version", "algorithm", "encrypted_payload", "encrypted_aes_key"};
            for (String field : requiredFields) {
                if (!packageJson.has(field)) {
                    throw new CompletionException(new ValueError("Missing required fields in encrypted package: " + field));
                }
            }

            // Validate version and algorithm
            String version = packageJson.get("version").getAsString();
            String algorithm = packageJson.get("algorithm").getAsString();

            if (!Constants.PROTOCOL_VERSION.equals(version)) {
                throw new CompletionException(new ValueError(
                    "Unsupported protocol version: '" + version + "'. Expected: '" + Constants.PROTOCOL_VERSION + "'"));
            }
            if (!Constants.HYBRID_ALGORITHM.equals(algorithm)) {
                throw new CompletionException(new ValueError(
                    "Unsupported encryption algorithm: '" + algorithm + "'. Expected: '" + Constants.HYBRID_ALGORITHM + "'"));
            }

            // Validate encrypted_payload structure
            JsonObject encryptedPayload = packageJson.get("encrypted_payload").getAsJsonObject();
            String[] payloadRequired = {"ciphertext", "nonce", "tag"};
            for (String field : payloadRequired) {
                if (!encryptedPayload.has(field)) {
                    throw new CompletionException(new ValueError("Missing fields in encrypted_payload: " + field));
                }
            }

            // Guard: private key must be initialized
            if (this.privateKey == null) {
                throw new CompletionException(new SecurityError("Private key not initialized. Call generateKeys() or loadKeys() first."));
            }

            try {
                // Decrypt AES key with private key
                byte[] encryptedAESKey = Base64.getDecoder().decode(packageJson.get("encrypted_aes_key").getAsString());
                
                Cipher rsaCipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
                rsaCipher.init(Cipher.DECRYPT_MODE, this.privateKey);
                byte[] aesKeyBytes = rsaCipher.doFinal(encryptedAESKey);
                SecretKeySpec aesKey = new SecretKeySpec(aesKeyBytes, "AES");

                // Decrypt payload
                byte[] ciphertext = Base64.getDecoder().decode(encryptedPayload.get("ciphertext").getAsString());
                byte[] nonce = Base64.getDecoder().decode(encryptedPayload.get("nonce").getAsString());
                byte[] tag = Base64.getDecoder().decode(encryptedPayload.get("tag").getAsString());

                Cipher aesCipher = Cipher.getInstance("AES/GCM/NoPadding");
                aesCipher.init(Cipher.DECRYPT_MODE, aesKey, new GCMParameterSpec(Constants.GCM_TAG_SIZE * 8, nonce));
                
                // Combine ciphertext (without tag) and tag for decryption
                byte[] ciphertextWithTag = new byte[ciphertext.length + tag.length];
                System.arraycopy(ciphertext, 0, ciphertextWithTag, 0, ciphertext.length);
                System.arraycopy(tag, 0, ciphertextWithTag, ciphertext.length, tag.length);

                byte[] plaintextBytes = aesCipher.doFinal(ciphertextWithTag);

                // Parse JSON response
                Map<String, Object> response;
                try {
                    Object parsed = gson.fromJson(new String(plaintextBytes, StandardCharsets.UTF_8), Object.class);
                    @SuppressWarnings("unchecked")
                    Map<String, Object> resultMap = (Map<String, Object>) parsed;
                    response = resultMap != null ? resultMap : new HashMap<>();
                } catch (Exception e) {
                    throw new CompletionException(new ValueError("Decrypted response is not valid JSON: " + e.getMessage()));
                }

                // Add metadata
                if (!response.containsKey("_metadata")) {
                    response.put("_metadata", new HashMap<String, Object>());
                }
                @SuppressWarnings("unchecked")
                Map<String, Object> metadata = (Map<String, Object>) response.get("_metadata");
                metadata.put("payload_id", payloadId);
                metadata.put("processed_at", packageJson.has("processed_at") ? packageJson.get("processed_at").getAsString() : null);
                metadata.put("is_encrypted", true);
                metadata.put("encryption_algorithm", Constants.HYBRID_ALGORITHM);

                return response;

            } catch (Exception e) {
                throw new CompletionException(new SecurityError("Decryption failed: integrity check or authentication failed"));
            }
        });
    }

    /**
     * Error class for invalid argument/value errors (maps to Python ValueError).
     */
    public static class ValueError extends Exception {
        public ValueError(String message) {
            super(message);
        }

        public ValueError(String message, Throwable cause) {
            super(message, cause);
        }
    }
}