/*
 * Copyright (c) 2016 NetApp
 * All rights reserved
 */

package com.netapp.oci.platform.common.interfaces.session;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.jboss.naming.remote.client.InitialContextFactory;
import org.wildfly.security.ssl.Protocol;
import org.xnio.Options;

import javax.ejb.CreateException;
import javax.ejb.EJBException;
import javax.naming.Context;
import javax.naming.InitialContext;
import javax.naming.NamingException;
import javax.net.ssl.*;
import java.io.IOException;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.rmi.RemoteException;
import java.rmi.UnmarshalException;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Properties;

public class RemoteAppSessionFactory {

    public static final RemoteAppSessionFactory INSTANCE = new RemoteAppSessionFactory();

    private static final Logger logger = LogManager.getLogger();

    private static final String CONNECT_OPTIONS_PREFIX = "jboss.naming.client.connect.options.";

    private InitialContext rootEjbNamingContext = null;
    private volatile boolean inited = false;
    private String hostname;
    private String port;
    private String user;
    private String password;
    private X509ExtendedKeyManager keyManager;
    private X509TrustManager trustManager;
    private int timeoutMillis;
    private boolean logExceptions;

    private RemoteAppSessionFactory() {
    }

    /**
     * Initialization of remote app session factory using password based authentication and an optional custom trustmanager
     *
     * @param hostname      the hostname to connect to (e.g. 127.0.0.1)
     * @param port          the port to connect to (e.g. 80 or 443)
     * @param user          the user for connecting to server (e.g. admin)
     * @param password      the password for connecting to server (e.g. well, I'm not telling you!)
     * @param trustManager  custom trustManager for use with a secure connection (SSL)
     */
    public synchronized void init(String hostname, String port, String user, String password, X509TrustManager trustManager) {
        this.init(hostname, port, user, password, trustManager, 1000, true);
    }

    /**
     * Initialization of remote app session factory using certificate based authentication
     *
     * @param hostname            the hostname to connect to (e.g. 127.0.0.1)
     * @param port                the port to connect to (e.g. 80 or 443)
     * @param trustManager        the trust manager for server authentication
     * @param keyManager          the key manager to load the client public key for client-cert authentication
     */
    public synchronized void init(String hostname, String port, X509TrustManager trustManager, X509ExtendedKeyManager keyManager) {
        this.init(hostname, port, trustManager, keyManager, 1000, true);
    }

    /**
     * Initialization of remote app session factory using password based authentication and an optional custom trustmanager
     *
     * @param hostname      the hostname to connect to (e.g. 127.0.0.1)
     * @param port          the port to connect to (e.g. 80 or 443)
     * @param user          the user for connecting to server (e.g. admin)
     * @param password      the password for connecting to server (e.g. well, I'm not telling you!)
     * @param trustManager  custom trustManager for use with a secure connection (SSL)
     * @param timeoutMillis the timeout for lookup calls
     * @param logExceptions true if exceptions are to be logged, false otherwise
     */
    public synchronized void init(String hostname, String port, String user, String password, X509TrustManager trustManager, int timeoutMillis, boolean logExceptions) {
        this.hostname = hostname;
        this.port = port;
        this.user = user;
        this.password = password;
        this.trustManager = trustManager;
        this.keyManager = null;
        this.timeoutMillis = timeoutMillis;
        this.logExceptions = logExceptions;
        this.inited = true;
    }

    /**
     * Initialization of remote app session factory using certificate based authentication
     *
     * @param hostname          the hostname to connect to (e.g. 127.0.0.1)
     * @param port              the port to connect to (e.g. 80 or 443)
     * @param trustManager      the trust manager class for server authentication
     * @param keyManager        the key manager class to load the client public key for client-cert authentication
     * @param timeoutMillis     the timeout for lookup calls
     * @param logExceptions     true if exceptions are to be logged, false otherwise
     */
    public synchronized void init(String hostname, String port, X509TrustManager trustManager, X509ExtendedKeyManager keyManager, int timeoutMillis, boolean logExceptions) {
        this.hostname = hostname;
        this.port = port;
        this.trustManager = trustManager;
        this.keyManager = keyManager;
        this.timeoutMillis = timeoutMillis;
        this.logExceptions = logExceptions;
        this.inited = true;
    }

    /**
     * Using WildFly Naming client library to create the root JNDI context.
     * See http://www.mastertheboss.com/jboss-server/jboss-as-7/jboss-as-7-remote-ejb-client-tutorial?showall=
     * For previous methods see: https://blog.akquinet.de/2014/09/26/jboss-eap-wildfly-three-ways-to-invoke-remote-ejbs/
     * </p>
     */
    private InitialContext createRootEjbNamingContext() throws NamingException {
        final Properties prop = new Properties();

        prop.put(Context.INITIAL_CONTEXT_FACTORY, "org.wildfly.naming.client.WildFlyInitialContextFactory");

        final String protocol = "remote+https";

        // EJB Client Global Properties
        // Boolean value that specifies whether the SSL protocol is enabled for all connections
        prop.put(CONNECT_OPTIONS_PREFIX + Options.SSL_ENABLED, "true");

        // EJB Client Connection Properties
        prop.put(CONNECT_OPTIONS_PREFIX + Options.SSL_STARTTLS, "true");

        if (keyManager != null || trustManager != null) {
            // WildFly 11 will ignore the legacy trustmanager and clientkeymanager properties. Instead,
            // WildFly will use the default context. So, we modify the SSL context default to be sure our custom managers are used.
            try {

                SSLContext context = SSLContext.getInstance(Protocol.TLSv1_2.name);
                context.init(
                        keyManager == null ? null : new KeyManager[]{keyManager},
                        trustManager == null ? null : new TrustManager[]{trustManager},
                        null);
                SSLContext.setDefault(context);
            } catch (NoSuchAlgorithmException | KeyManagementException e) {
                logger.error("Failed to initialize remote app, error message: " + e.getMessage());
            }
        }

        // tests for instance set the port to "-1" as they set the URL to "https://localhost".
        // figure out actual connection port from secure status
        String connectionPort = port;
        if ("-1".equals(port)) {
            connectionPort = "443";
        }
        prop.put(Context.PROVIDER_URL, protocol + "://" + hostname + ":" + connectionPort);

        // EJB Client Connection Properties
        if (user != null && password != null) {
            prop.put("java.naming.security.principal", user);
            String passwordBase64 = Base64.getEncoder().encodeToString(password.getBytes());
            //noinspection deprecation
            prop.put(InitialContextFactory.PASSWORD_BASE64_KEY , passwordBase64);
        }

        // Boolean value that determines whether credentials must be provided by the client to connect successfully. The default value is true.
        // If set to true, the client must provide credentials. If set to false, invocation is allowed as long as the remoting connector does not request a security realm.
        prop.put(CONNECT_OPTIONS_PREFIX + Options.SASL_POLICY_NOANONYMOUS, "false");
        // Boolean value that enables or disables the use of plain text messages during the authentication.
        // If using JAAS, it must be set to false to allow a plain text password.
        prop.put(CONNECT_OPTIONS_PREFIX + Options.SASL_POLICY_NOPLAINTEXT, "false");

        return new InitialContext(prop);
    }

    /**
     * Look up the session for the specified JNDI name. The returned proxy will not handle tries.
     *
     * @return A proxy of the session that won't handle retries.
     * @throws NamingException if failed to look up the session
     */
    public synchronized <E> E lookupNoRetry(final String jndiName) throws NamingException {
        return this.lookup(jndiName, 0);
    }

    public synchronized <E> E lookup(final String jndiName) throws NamingException {
        return this.lookup(jndiName, timeoutMillis);
    }

    private <E> E lookup(final String jndiName, int timeoutMillis) throws NamingException {
        if (!isInited())
            throw new IllegalStateException("Need to init");
        final String resolvedJndiName = resolveJndiName(jndiName);
        final Object session = lookupSession(resolvedJndiName);
        //noinspection unchecked
        return (E) Proxy.newProxyInstance(
                RemoteAppSessionFactory.class.getClassLoader(),
                session.getClass().getInterfaces(),
                new RemoteAppSessionFactory.ProxyInvocationHandler(resolvedJndiName, timeoutMillis));
    }

    private synchronized Object lookupSession(String resolvedJndiName) throws NamingException {
        if (!isInited()) {
            throw new IllegalStateException("Communication Manager is currently stopping!");
        }
        if (rootEjbNamingContext == null) {
            rootEjbNamingContext = createRootEjbNamingContext();
        }
        return rootEjbNamingContext.lookup(resolvedJndiName);
    }

    private String resolveJndiName(String jndiName) {
        String ejbJNDIName;
        if (jndiName.startsWith("java:global/")) {
            if (isNullApp(jndiName)) {
                // preserve the slash after java:global so it looks like /(module)
                ejbJNDIName = jndiName.replaceFirst("java:global/", "/");
            } else {
                // application name does not need a slash: e.g: (app)/(module)
                ejbJNDIName = jndiName.replaceFirst("java:global/", "");
            }
        } else ejbJNDIName = jndiName;
        return ejbJNDIName;
    }

    private boolean isNullApp(String jndiName) {
        int slashes = 0;
        for (int i = 0; i < jndiName.length(); i++) {
            char c = jndiName.charAt(i);
            if (c == '/') {
                slashes++;
            }
            if (c == '!') {
                break;
            }
        }
        return slashes == 2;
    }

    public synchronized boolean isInited() { return inited; }

    public synchronized void close() {
        closeRootEjbNamingContext();
        inited = false;
    }

    private synchronized void closeRootEjbNamingContext() {
        if (rootEjbNamingContext != null) {
            try {
                rootEjbNamingContext.close();
            } catch (NamingException e) {
                logger.error("Failed to close EjbNamingContext", e);
            } finally {
                rootEjbNamingContext = null;
            }
        }
    }

    /**
     * Represent a wrapper or handler class to the underline EJB Session.
     * The main intention of this class is to create a generic way to encapsulate problem related to:
     * <p/>
     * <ul>
     * <li> Create an instance of the right type.
     * <li> Provides a mechanism of reconnection is case of a broken communication.
     * </ul>
     * <p/>
     * Once established a session the User can call EJB session's method that are
     * executed by {@link ProxyInvocationHandler#invoke(Object, Method, Object[])}
     * which uses underling implementation represented by <code>obj</code>.
     * <code>obj</code> is an instance of a particular EJB Session.  In order to create the right EJB Session
     * we use the <code>jndiName</code>
     */
    private class ProxyInvocationHandler implements InvocationHandler {
        private String jndiName;
        private int timeoutMillis;

        public ProxyInvocationHandler(final String jndiName, int timeoutMillis) {
            this.jndiName = jndiName;
            this.timeoutMillis = timeoutMillis;
        }

        public Object invoke(Object proxy, Method m, Object[] args) throws Throwable {
            final int retries = Math.max(timeoutMillis / 1000, 1);
            final int sleepTime = Math.max(timeoutMillis / retries, 1000);

            int retriesLeft = retries;
            while (isInited()) {
                try {
                    return m.invoke(lookupSession(jndiName), args);
                } catch (InvocationTargetException e) {
                    final Throwable targetException = e.getTargetException();

                    if (targetException == null) {
                        throw e;
                    }

                    if (timeoutMillis > 0 && !isRecoverableCommunicationError(targetException)) {
                        if (logExceptions) {
                            logger.error("RemoteAppSessionFactory - Failed to communicate with the server ( " + jndiName + ") - unrecoverable error: " + targetException.getMessage(), targetException);
                        }
                        throw targetException;
                    }

                    // We are handling only these kind of exceptions because this means that the target method that we're executing
                    // threw an exception. This means that EJB communication layer failed somehow when calling the server.
                    // Stop and Restart, this will create and initialize all the related instances.
                    retriesLeft--;

                    if (retriesLeft <= 0) {
                        if (logExceptions) {
                            logger.error("RemoteAppSessionFactory - Failed to communicate with the server (" + jndiName + ") - no retries left: " + targetException.getMessage(), targetException);
                        }
                        throw targetException;
                    }

                    // Provide only basic info in the log. The detailed exception, including stack trace will be written when
                    // no retries are left.
                    if (logExceptions) {
                        logger.warn("RemoteAppSessionFactory - Failed to communicate with the server (" + jndiName + ") - " + retriesLeft + " retries left: " + targetException.getClass().getName() + ": " + targetException.getMessage());
                    }

                    try {
                        Thread.sleep(sleepTime);
                    } catch (InterruptedException ie) {
                        // Fall through...
                    }
                }
            }

            // If we're here, the service is stopping...
            throw new InterruptedException("Session was stopped");
        }

        /**
         * HACK:
         * On these exceptions, we will try to silently retry the connection before propagating the problem
         */
        private boolean isRecoverableCommunicationError(Throwable targetException) {
            // HACK: If the server sent us an exception and we could not marshal this exception, there is no point in
            // recovering
            try {
                if (targetException instanceof UnmarshalException) {
                    if (targetException.getCause() instanceof ClassNotFoundException)
                        return false;
                }
            } catch (NoClassDefFoundError e) {
                return false;
            }

            // Exceptions thrown by the JBoss communication layer
            if (targetException instanceof RemoteException)
                return true;

            // Exceptions thrown during session.create
            if (targetException instanceof CreateException)
                return true;

            // HACK: Exception usually thrown when the server reboots and the client was in the middle of an RMI call
            if (targetException.getClass().getName().equals(IllegalArgumentException.class.getName()))
                return true;

            // OutOfMemory is also a recoverable error...
            if (targetException instanceof OutOfMemoryError)
                return true;

            if (targetException instanceof EJBException) {
                // EJB Channel might be broken, should be recoverable
                if (targetException.getCause() instanceof IOException)
                    return true;
            }

            return false;
        }
    }
}