package me.proxylink.plugin.tunnel;

import me.proxylink.common.AgentConfig;

import javax.net.ServerSocketFactory;
import javax.net.SocketFactory;
import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.util.function.Consumer;

final class TlsSocketFactory {
    private TlsSocketFactory() {
    }

    static ServerSocket createServerSocket(AgentConfig config) throws IOException, GeneralSecurityException {
        InetSocketAddress bindAddress = new InetSocketAddress(
                config.tunnel().listenHost(),
                config.tunnel().listenPort()
        );

        if (!config.tls().enabled()) {
            ServerSocket serverSocket = new ServerSocket();
            serverSocket.setReuseAddress(true);
            serverSocket.bind(bindAddress);
            return serverSocket;
        }

        SSLContext context = sslContext(
                config.tls().keyStorePath(),
                config.tls().keyStorePassword(),
                config.tls().trustStorePath(),
                config.tls().trustStorePassword(),
                null
        );
        ServerSocketFactory factory = context.getServerSocketFactory();
        SSLServerSocket serverSocket = (SSLServerSocket) factory.createServerSocket();
        serverSocket.setReuseAddress(true);
        serverSocket.setNeedClientAuth(config.tls().requireClientAuth());
        serverSocket.setEnabledProtocols(enabledProtocols(serverSocket.getSupportedProtocols()));
        serverSocket.bind(bindAddress);
        return serverSocket;
    }

    static Socket createClientSocket(AgentConfig config) throws IOException, GeneralSecurityException {
        return createClientSocket(config, ignored -> {
        });
    }

    static Socket createClientSocket(AgentConfig config, Consumer<String> learnedFingerprintConsumer)
            throws IOException, GeneralSecurityException {
        AgentConfig.FrontendEndpointConfig endpoint = new AgentConfig.FrontendEndpointConfig(
                "default",
                config.tunnel().connectHost(),
                config.tunnel().connectPort(),
                config.tls().pinnedCertificateSha256()
        );
        return createClientSocket(config, endpoint, learnedFingerprintConsumer);
    }

    static Socket createClientSocket(
            AgentConfig config,
            AgentConfig.FrontendEndpointConfig endpoint,
            Consumer<String> learnedFingerprintConsumer
    ) throws IOException, GeneralSecurityException {
        Socket socket;
        PinnedCertificateTrustManager pinningTrustManager = null;

        if (!config.tls().enabled()) {
            socket = SocketFactory.getDefault().createSocket();
        } else {
            pinningTrustManager = pinningTrustManager(config, endpoint);
            SSLContext context = sslContext(
                    config.tls().keyStorePath(),
                    config.tls().keyStorePassword(),
                    config.tls().trustStorePath(),
                    config.tls().trustStorePassword(),
                    pinningTrustManager
            );
            socket = context.getSocketFactory().createSocket();
        }

        socket.setTcpNoDelay(true);
        socket.setKeepAlive(true);
        socket.connect(
                new InetSocketAddress(endpoint.connectHost(), endpoint.connectPort()),
                config.tunnel().connectTimeoutMillis()
        );

        if (socket instanceof SSLSocket sslSocket) {
            sslSocket.setEnabledProtocols(enabledProtocols(sslSocket.getSupportedProtocols()));
            sslSocket.startHandshake();
            if (pinningTrustManager != null && pinningTrustManager.learnedFingerprint() != null) {
                learnedFingerprintConsumer.accept(pinningTrustManager.learnedFingerprint());
            }
        }

        return socket;
    }

    private static SSLContext sslContext(
            Path keyStorePath,
            String keyStorePassword,
            Path trustStorePath,
            String trustStorePassword,
            TrustManager pinningTrustManager
    ) throws IOException, GeneralSecurityException {
        KeyManager[] keyManagers = keyStorePath == null ? null : keyManagers(keyStorePath, keyStorePassword);
        TrustManager[] trustManagers;
        if (pinningTrustManager != null) {
            trustManagers = new TrustManager[]{pinningTrustManager};
        } else {
            trustManagers = trustStorePath == null ? null : trustManagers(trustStorePath, trustStorePassword);
        }

        SSLContext context = SSLContext.getInstance("TLS");
        context.init(keyManagers, trustManagers, null);
        return context;
    }

    private static PinnedCertificateTrustManager pinningTrustManager(
            AgentConfig config,
            AgentConfig.FrontendEndpointConfig endpoint
    ) {
        String pinnedCertificate = endpoint.pinnedCertificateSha256().isBlank()
                ? config.tls().pinnedCertificateSha256()
                : endpoint.pinnedCertificateSha256();
        if (!pinnedCertificate.isBlank() || config.tls().trustOnFirstUse()) {
            return new PinnedCertificateTrustManager(
                    pinnedCertificate,
                    config.tls().trustOnFirstUse()
            );
        }
        return null;
    }

    private static KeyManager[] keyManagers(Path keyStorePath, String password)
            throws IOException, GeneralSecurityException {
        KeyStore keyStore = loadStore(keyStorePath, password);
        KeyManagerFactory factory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        factory.init(keyStore, password.toCharArray());
        return factory.getKeyManagers();
    }

    private static TrustManager[] trustManagers(Path trustStorePath, String password)
            throws IOException, GeneralSecurityException {
        KeyStore trustStore = loadStore(trustStorePath, password);
        TrustManagerFactory factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
        factory.init(trustStore);
        return factory.getTrustManagers();
    }

    private static KeyStore loadStore(Path path, String password) throws IOException, GeneralSecurityException {
        KeyStore store = KeyStore.getInstance(KeyStore.getDefaultType());
        try (InputStream input = Files.newInputStream(path)) {
            store.load(input, password.toCharArray());
        }
        return store;
    }

    private static String[] enabledProtocols(String[] supported) {
        boolean tls13 = false;
        boolean tls12 = false;
        for (String protocol : supported) {
            if ("TLSv1.3".equals(protocol)) {
                tls13 = true;
            }
            if ("TLSv1.2".equals(protocol)) {
                tls12 = true;
            }
        }
        if (tls13 && tls12) {
            return new String[]{"TLSv1.3", "TLSv1.2"};
        }
        if (tls13) {
            return new String[]{"TLSv1.3"};
        }
        if (tls12) {
            return new String[]{"TLSv1.2"};
        }
        return supported;
    }
}
