package com.cygnus.cloud.security;

import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.concurrent.ConcurrentHashMap;

final class PemKeyLoader {

    private static final ConcurrentHashMap<String, RSAPrivateKey> PRIVATE_KEYS =
            new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, RSAPublicKey> PUBLIC_KEYS =
            new ConcurrentHashMap<>();

    private PemKeyLoader() {
    }

    static RSAPrivateKey privateKey(String location) {
        return PRIVATE_KEYS.computeIfAbsent(location, PemKeyLoader::loadPrivateKey);
    }

    private static RSAPrivateKey loadPrivateKey(String location) {
        try {
            String encoded = read(location)
                    .replace("-----BEGIN PRIVATE KEY-----", "")
                    .replace("-----END PRIVATE KEY-----", "")
                    .replaceAll("\\s", "");
            PrivateKey key = KeyFactory.getInstance("RSA")
                    .generatePrivate(new PKCS8EncodedKeySpec(
                            Base64.getDecoder().decode(encoded)));
            return (RSAPrivateKey) key;
        } catch (Exception exception) {
            throw new IllegalStateException("Unable to load RSA private key", exception);
        }
    }

    static RSAPublicKey publicKey(String location) {
        return PUBLIC_KEYS.computeIfAbsent(location, PemKeyLoader::loadPublicKey);
    }

    private static RSAPublicKey loadPublicKey(String location) {
        try {
            String encoded = read(location)
                    .replace("-----BEGIN PUBLIC KEY-----", "")
                    .replace("-----END PUBLIC KEY-----", "")
                    .replaceAll("\\s", "");
            PublicKey key = KeyFactory.getInstance("RSA")
                    .generatePublic(new X509EncodedKeySpec(
                            Base64.getDecoder().decode(encoded)));
            return (RSAPublicKey) key;
        } catch (Exception exception) {
            throw new IllegalStateException("Unable to load RSA public key", exception);
        }
    }

    private static String read(String location) throws Exception {
        if (location == null || location.isBlank()) {
            throw new IllegalArgumentException("RSA key location is not configured");
        }
        if (location.contains("-----BEGIN ")) {
            return location;
        }
        if (location.startsWith("classpath:")) {
            String resource = location.substring("classpath:".length());
            try (InputStream stream = Thread.currentThread()
                    .getContextClassLoader()
                    .getResourceAsStream(resource)) {
                if (stream == null) {
                    throw new IllegalArgumentException("Key resource not found");
                }
                return new String(stream.readAllBytes(), StandardCharsets.US_ASCII);
            }
        }
        String file = location.startsWith("file:") ? location.substring(5) : location;
        return Files.readString(Path.of(file), StandardCharsets.US_ASCII);
    }
}