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 javax.crypto.*;
import java.io.*;
import java.math.BigInteger;
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.PKCS8EncodedKeySpec;
import java.security.spec.RSAKeyGenParameterSpec;
import java.io.File;
import java.io.UnsupportedEncodingException;
import java.security.spec.RSAPrivateCrtKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.util.Arrays;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Low-level secure completion client for the NOMYO API.
 *
 * <p>This class handles key management, hybrid encryption, HTTP communication
 * with retry logic, and response decryption. It is the core of the NOMYO
 * Java client and is used internally by {@link SecureChatCompletion}.</p>
 *
 * <h3>Encryption Wire Format</h3>
 * <p>Encrypted payloads use hybrid encryption (AES-256-GCM + RSA-4096-OAEP-SHA256):</p>
 * <ul>
 *   <li>A per-request 256-bit AES key encrypts the payload via AES-256-GCM</li>
 *   <li>The AES key is encrypted via RSA-4096-OAEP-SHA256 using the server's public key</li>
 *   <li>The result is a JSON package with base64-encoded ciphertext, nonce, tag, and encrypted AES key</li>
 * </ul>
 *
 * <h3>Key Lifecycle</h3>
 * <p>Client RSA keys are generated on first use (if not loaded from disk) and
 * reused across all subsequent calls until the client is discarded. Keys can
 * be persisted to disk via {@link #generateKeys(boolean, String, String)} or
 * loaded from disk via {@link #loadKeys(String, String, String)}.</p>
 */
public class SecureCompletionClient {

    // ── Instance Attributes ─────────────────────────────────────────

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

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

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

    /**
     * Maximum number of retries for retryable errors.
     */
    @Getter
    private final int maxRetries;

    /**
     * Whether secure memory operations are active.
     */
    @Getter
    private final boolean useSecureMemory;

    /**
     * Lock for double-checked key initialization.
     */
    private final ReentrantLock keyInitLock = new ReentrantLock();

    /**
     * RSA private key, or {@code null} if not yet loaded/generated.
     */
    @Getter
    private PrivateKey privateKey;

    // ── Internal State ──────────────────────────────────────────────

    /**
     * PEM-encoded public key string, or {@code null} if not yet loaded/generated.
     */
    @Getter
    private String publicPemKey;

    /**
     * Whether keys have been initialized.
     */
    private volatile boolean keysInitialized = false;

    /**
     * Constructs a {@code SecureCompletionClient} with default settings.
     */
    public SecureCompletionClient() {
        this(Constants.DEFAULT_BASE_URL, false, true, Constants.DEFAULT_MAX_RETRIES);
    }

    /**
     * Constructs a {@code SecureCompletionClient} with the specified settings.
     *
     * @param routerUrl    the NOMYO router base URL
     * @param allowHttp    whether to permit HTTP (non-HTTPS) URLs
     * @param secureMemory whether to enable memory locking/zeroing
     * @param maxRetries   number of 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;
    }

    // ── Key Management ──────────────────────────────────────────────

    private static String getEncryptedPrivateKeyFromFile(String privateKeyPath) {
        File myObj = new File(privateKeyPath);

        StringBuilder builder = new StringBuilder();

        try (Scanner myReader = new Scanner(myObj)) {
            while (myReader.hasNextLine()) {
                builder.append(myReader.nextLine());
            }
        } catch (FileNotFoundException e) {
            throw new RuntimeException("Tried to load private key from disk but no file found" + e.getMessage());
        }

        return builder.toString();
    }

    /**
     * Generates a new 4096-bit RSA key pair.
     *
     * <p>The public exponent is fixed at 65537. The generated key pair
     * is stored in memory. Use {@code saveToDir} to persist to disk.</p>
     *
     * @param saveToFile whether to save the keys to disk
     * @param keyDir     directory to save keys (ignored if {@code saveToFile} is {@code false})
     * @param password   optional password to encrypt the private key file
     */
    public void generateKeys(boolean saveToFile, String keyDir, String password) {
        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) {
                File keyFolder = new File(keyDir);
                if (!keyFolder.exists() && !keyFolder.mkdirs()) {
                    throw new IOException("Failed to create key directory: " + keyDir);
                }

                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 (FileWriter fileWriter = new FileWriter(privateKeyPath.toFile())) {
                        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 e) {
                                throw new RuntimeException(e);
                            }

                        }

                        fileWriter.write(privatePem);
                        fileWriter.flush();
                    }
                }

                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 (FileWriter fileWriter = new FileWriter(publicKeyPath.toFile())) {
                        fileWriter.write(publicPem);
                        fileWriter.flush();
                    }
                }
            }

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

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

    /**
     * Generates a new 4096-bit RSA key pair and saves to the default directory.
     *
     * @param saveToFile whether to save the keys to disk
     */
    public void generateKeys(boolean saveToFile) {
        generateKeys(saveToFile, Constants.DEFAULT_KEY_DIR, null);
    }

    /**
     * Loads an RSA private key from disk.
     *
     * <p>If {@code publicPemKeyPath} is {@code null}, the public key is
     * derived from the loaded private key. Validates that the key size
     * is at least {@link Constants#MIN_RSA_KEY_SIZE} bits.</p>
     *
     * @param privateKeyPath   path to the private key PEM file
     * @param publicPemKeyPath optional path to the public key PEM file
     * @param password         optional password for the encrypted private key
     */
    public void loadKeys(String privateKeyPath, String publicPemKeyPath, String password) {
        File keyFile = new File(privateKeyPath);
        if (!keyFile.exists()) {
            throw new RuntimeException("Private key file not found: " + privateKeyPath);
        }

        String keyContent;
        if (password != null && !password.isEmpty()) {
            keyContent = getEncryptedPrivateKeyFromFile(privateKeyPath);

            try {
                keyContent = Pass2Key.decrypt("AES/GCM/NoPadding", keyContent, password);
            } catch (NoSuchPaddingException | NoSuchAlgorithmException
                     | BadPaddingException | IllegalBlockSizeException | InvalidAlgorithmParameterException |
                     InvalidKeyException e) {
                System.out.println("Wrong password!");
                return;
            }
        } else {
            keyContent = getEncryptedPrivateKeyFromFile(privateKeyPath);
        }

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

    /**
     * Loads an RSA private key from disk, deriving the public key.
     *
     * @param privateKeyPath path to the private key PEM file
     * @param password       optional password for the encrypted private key
     */
    public void loadKeys(String privateKeyPath, String password) {
        loadKeys(privateKeyPath, null, password);
    }

    // ── Server Key Fetching ─────────────────────────────────────────

    /**
     * Fetches the server's RSA public key from the PKI endpoint.
     *
     * <p>Performs a GET request to {@code {routerUrl}/pki/public_key}
     * and returns the server PEM public key as a string. Validates that
     * the response parses as a valid PEM public key.</p>
     *
     * @return the server's PEM-encoded public key string
     * @throws SecurityError if the URL is not HTTPS and {@code allowHttp} is {@code false},
     *                       or if the response does not contain a valid PEM public key
     */
    public CompletableFuture<String> fetchServerPublicKey() {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    // ── Encryption ──────────────────────────────────────────────────

    /**
     * Encrypts a payload dict using hybrid encryption.
     *
     * <p>Serializes the payload to JSON, then encrypts it using:
     * <ol>
     *   <li>A per-request 256-bit AES key (AES-256-GCM)</li>
     *   <li>RSA-OAEP-SHA256 wrapping of the AES key with the server's public key</li>
     * </ol>
     * </p>
     *
     * @param payload the payload to encrypt (OpenAI-compatible chat parameters)
     * @return raw encrypted bytes (JSON package serialized to bytes)
     * @throws SecurityError if encryption fails or keys are not loaded
     */
    public CompletableFuture<byte[]> encryptPayload(Map<String, Object> payload) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Core hybrid encryption routine.
     */
    public CompletableFuture<byte[]> doEncrypt(byte[] payloadBytes, byte[] aesKey) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    // ── Decryption ──────────────────────────────────────────────────

    /**
     * Decrypts a server response.
     */
    public CompletableFuture<Map<String, Object>> decryptResponse(byte[] encryptedResponse, String payloadId) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    // ── Secure Request Lifecycle ────────────────────────────────────

    /**
     * Full request lifecycle: encrypt → HTTP POST → retry → decrypt → return.
     *
     * <h3>Request Headers</h3>
     * <pre>
     * Content-Type: application/octet-stream
     * X-Payload-ID: {payloadId}
     * X-Public-Key: {urlEncodedPublicPemKey}
     * Authorization: Bearer {apiKey}            (if apiKey is provided)
     * X-Security-Tier: {tier}                   (if securityTier is provided)
     * </pre>
     *
     * <h3>POST</h3>
     * {@code {routerUrl}/v1/chat/secure_completion} with encrypted payload as body.
     *
     * <h3>Retry Logic</h3>
     * <ul>
     *   <li>Retryable status codes: {@code {429, 500, 502, 503, 504}}</li>
     *   <li>Backoff: {@code 2^(attempt-1)} seconds (1s, 2s, 4s…)</li>
     *   <li>Total attempts: {@code maxRetries + 1}</li>
     *   <li>Network errors also retry</li>
     *   <li>Non-retryable exceptions propagate immediately</li>
     * </ul>
     *
     * <h3>Status → Exception Mapping</h3>
     * <table>
     *   <tr><th>Status</th><th>Result</th></tr>
     *   <tr><td>200</td><td>Return decrypted response map</td></tr>
     *   <tr><td>400</td><td>{@code InvalidRequestError}</td></tr>
     *   <tr><td>401</td><td>{@code AuthenticationError}</td></tr>
     *   <tr><td>403</td><td>{@code ForbiddenError}</td></tr>
     *   <tr><td>404</td><td>{@code APIError}</td></tr>
     *   <tr><td>429</td><td>{@code RateLimitError}</td></tr>
     *   <tr><td>500</td><td>{@code ServerError}</td></tr>
     *   <tr><td>503</td><td>{@code ServiceUnavailableError}</td></tr>
     *   <tr><td>502/504</td><td>{@code APIError} (retryable)</td></tr>
     *   <tr><td>other</td><td>{@code APIError} (non-retryable)</td></tr>
     *   <tr><td>network error</td><td>{@code APIConnectionError}</td></tr>
     * </table>
     *
     * @param payload      the payload to send (OpenAI-compatible chat parameters)
     * @param payloadId    unique payload identifier
     * @param apiKey       optional API key for authentication
     * @param securityTier optional security tier ({@code "standard"}, {@code "high"}, or {@code "maximum"})
     * @return the decrypted response map
     * @throws SecurityError           if encryption/decryption fails
     * @throws APIConnectionError      if a network error occurs
     * @throws InvalidRequestError     if the API returns 400
     * @throws AuthenticationError     if the API returns 401
     * @throws ForbiddenError          if the API returns 403
     * @throws RateLimitError          if the API returns 429
     * @throws ServerError             if the API returns 500
     * @throws ServiceUnavailableError if the API returns 503
     * @throws APIError                for other non-retryable errors
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(
            Map<String, Object> payload,
            String payloadId,
            String apiKey,
            String securityTier
    ) {
        throw new UnsupportedOperationException("Not yet implemented");
    }

    /**
     * Sends a secure request without a security tier.
     *
     * @param payload   the payload to send
     * @param payloadId unique payload identifier
     * @param apiKey    optional API key for authentication
     * @return the decrypted response map
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(
            Map<String, Object> payload,
            String payloadId,
            String apiKey
    ) {
        return sendSecureRequest(payload, payloadId, apiKey, null);
    }

    /**
     * Sends a secure request with no API key or security tier.
     *
     * @param payload   the payload to send
     * @param payloadId unique payload identifier
     * @return the decrypted response map
     */
    public CompletableFuture<Map<String, Object>> sendSecureRequest(
            Map<String, Object> payload,
            String payloadId
    ) {
        return sendSecureRequest(payload, payloadId, null, null);
    }

    // ── Key Initialization ──────────────────────────────────────────

    /**
     * Ensures RSA keys are loaded or generated.
     *
     * <p>Uses double-checked locking via {@link ReentrantLock} to ensure
     * thread-safe initialization. If {@code keyDir} is set, attempts to
     * load keys from disk first; if that fails, generates new keys.</p>
     *
     * @param keyDir directory to persist keys, or {@code null} for ephemeral
     */
    public void ensureKeys(String keyDir) {
        if (keysInitialized) return;
        keyInitLock.lock();
        try {
            if (keysInitialized) return;
            // TODO: implement key loading/generation
            keysInitialized = true;
        } finally {
            keyInitLock.unlock();
        }
    }

    // ── Key Validation ──────────────────────────────────────────────

    /**
     * Validates that an RSA key meets the minimum size requirement.
     *
     * @param key the RSA key to validate
     * @throws SecurityError if the key size is less than {@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 {
            java.security.KeyFactory kf = java.security.KeyFactory.getInstance("RSA");
            java.security.spec.RSAPrivateCrtKeySpec crtSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateCrtKeySpec.class);
            return crtSpec.getModulus().bitLength();
        } catch (Exception ignored) {
            // Try RSAPrivateKeySpec
            try {
                java.security.KeyFactory kf = java.security.KeyFactory.getInstance("RSA");
                java.security.spec.RSAPrivateKeySpec privSpec = kf.getKeySpec(key, java.security.spec.RSAPrivateKeySpec.class);
                return privSpec.getModulus().bitLength();
            } catch (Exception ignored2) {
                // Fall back to encoded length
                if (key.getEncoded() != null) {
                    return key.getEncoded().length * 8;
                }
            }
        }
        return 0;
    }

    // ── HTTP Status → Exception Mapping ─────────────────────────────

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

    // ── URL Encoding ────────────────────────────────────────────────

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

    // ── Getters ─────────────────────────────────────────────────────

    /**
     * Closes the client and releases any resources.
     */
    public void close() {
        throw new UnsupportedOperationException("Not yet implemented");
    }
}