/*
 * Copyright The WildFly Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package org.jboss.as.server.mgmt.domain;

import static org.wildfly.common.Assert.checkNotNullParam;
import java.io.DataInput;
import java.io.IOException;
import java.io.InterruptedIOException;
import java.net.URI;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import org.jboss.as.controller.ModelController;
import org.jboss.as.controller.remote.ResponseAttachmentInputStreamSupport;
import org.jboss.as.controller.remote.TransactionalProtocolClient;
import org.jboss.as.controller.remote.TransactionalProtocolOperationHandler;
import org.jboss.as.protocol.ProtocolConnectionConfiguration;
import org.jboss.as.protocol.ProtocolConnectionManager;
import org.jboss.as.protocol.ProtocolConnectionUtils;
import org.jboss.as.protocol.StreamUtils;
import org.jboss.as.protocol.mgmt.AbstractManagementRequest;
import org.jboss.as.protocol.mgmt.ActiveOperation;
import org.jboss.as.protocol.mgmt.FlushableDataOutput;
import org.jboss.as.protocol.mgmt.FutureManagementChannel;
import org.jboss.as.protocol.mgmt.ManagementChannelHandler;
import org.jboss.as.protocol.mgmt.ManagementPingRequest;
import org.jboss.as.protocol.mgmt.ManagementRequestContext;
import org.jboss.as.remoting.management.ManagementRemotingServices;
import org.jboss.as.server.logging.ServerLogger;
import org.jboss.dmr.ModelNode;
import org.jboss.remoting3.Channel;
import org.jboss.remoting3.Connection;
import org.wildfly.security.auth.client.AuthenticationContext;
import org.wildfly.security.manager.WildFlySecurityManager;

/**
 * The connection to the host-controller. In case the channel is closed it's the host-controllers responsibility
 * to ask individual managed servers to reconnect.
 *
 * @author Emanuel Muckenhuber
 */
class HostControllerConnection extends FutureManagementChannel {

    private static final String SERVER_CHANNEL_TYPE = ManagementRemotingServices.SERVER_CHANNEL;
    private static final long reconnectionDelay;

    static {
        // Since there is the remoting connection timeout we might not need a delay between reconnection attempts at all
        reconnectionDelay = Long.parseLong(WildFlySecurityManager.getPropertyPrivileged("jboss.as.domain.host.reconnection.delay", "1500"));
    }

    private final String serverProcessName;
    private final ProtocolConnectionManager connectionManager;
    private final ManagementChannelHandler channelHandler;
    private final ExecutorService executorService;
    private final int initialOperationID;
    private final ResponseAttachmentInputStreamSupport responseAttachmentSupport;

    private volatile ProtocolConnectionConfiguration configuration;
    private volatile AuthenticationContext authenticationContext;
    private volatile ReconnectRunner reconnectRunner;

    HostControllerConnection(final String serverProcessName, final int initialOperationID,
                             final ProtocolConnectionConfiguration configuration,
                             final AuthenticationContext authenticationContext,
                             final ResponseAttachmentInputStreamSupport responseAttachmentSupport,
                             final ExecutorService executorService) {
        this.serverProcessName = serverProcessName;
        this.configuration = configuration;
        this.authenticationContext = checkNotNullParam("authenticationContext", authenticationContext);
        this.initialOperationID = initialOperationID;
        this.executorService = executorService;
        this.channelHandler = new ManagementChannelHandler(this, executorService);
        this.connectionManager = ProtocolConnectionManager.create(configuration, this, new ReconnectTask());
        this.responseAttachmentSupport = responseAttachmentSupport;
    }

    ManagementChannelHandler getChannelHandler() {
        return channelHandler;
    }

    @Override
    public Channel getChannel() throws IOException {
        final Channel channel = super.getChannel();
        if(channel == null) {
            // Fail fast, don't try to await a new channel
            throw channelClosed();
        }
        return channel;
    }

    /**
     * Connect to the HC and retrieve the current model updates.
     *
     * @param controller the server controller
     * @param callback the operation completed callback
     *
     * @throws IOException for any error
     */
    synchronized void openConnection(final ModelController controller, final ActiveOperation.CompletedCallback<ModelNode> callback) throws Exception {
        boolean ok = false;
        final Connection connection = internalConnect();
        try {
            channelHandler.executeRequest(new ServerRegisterRequest(), null, callback);
            // HC is the same version, so it will support sending the subject
            channelHandler.getAttachments().attach(TransactionalProtocolClient.SEND_IDENTITY, Boolean.TRUE);
            channelHandler.getAttachments().attach(TransactionalProtocolClient.SEND_IN_VM, Boolean.TRUE);
            channelHandler.addHandlerFactory(new TransactionalProtocolOperationHandler(controller, channelHandler, responseAttachmentSupport));
            ok = true;
        } finally {
            if(!ok) {
                connection.close();
            }
        }
    }

    private Connection internalConnect() throws IOException {
        try {
            return authenticationContext.run(new PrivilegedExceptionAction<Connection>() {

                @Override
                public Connection run() throws Exception {
                    return connectionManager.connect();
                }
            });
        } catch (PrivilegedActionException e) {
            if (e.getCause() instanceof IOException) {
                throw (IOException) e.getCause();
            }

            throw new IllegalStateException(e);
        }
    }

    /**
     * This continuously tries to reconnect in a separate thread and will only stop if the connection was established
     * successfully or the server gets shutdown. If there is currently a reconnect task active the connection paramaters
     * and callback will get updated.
     *
     * @param reconnectUri    the updated connection uri
     * @param serverAuthToken the updated authentication token
     * @param callback        the current callback
     */
    synchronized void asyncReconnect(final URI reconnectUri, AuthenticationContext authenticationContext, final ReconnectCallback callback) {
        if (getState() != State.OPEN) {
            return;
        }

        // Update the configuration with the new URI
        final ProtocolConnectionConfiguration config = ProtocolConnectionConfiguration.copy(configuration);
        config.setUri(reconnectUri);
        this.configuration = config;
        // Update the authentication context with the new credentials
        this.authenticationContext = authenticationContext;

        final ReconnectRunner reconnectTask = this.reconnectRunner;
        if (reconnectTask == null) {
            final ReconnectRunner task = new ReconnectRunner();
            task.callback = callback;
            task.future = executorService.submit(task);
        } else {
            reconnectTask.callback = callback;
        }
    }

    /**
     * Reconnect to the HC.
     *
     * @return whether the server is still in sync
     * @throws IOException
     */
    synchronized boolean doReConnect() throws IOException {

        // In case we are still connected, test the connection and see if we can reuse it
        if(connectionManager.isConnected()) {
            try {
                final Future<Long> result = channelHandler.executeRequest(ManagementPingRequest.INSTANCE, null).getResult();
                result.get(15, TimeUnit.SECONDS); // Hmm, perhaps 15 is already too much
                return true;
            } catch (Exception e) {
                ServerLogger.AS_ROOT_LOGGER.debugf(e, "failed to ping the host-controller, going to reconnect");
            }
            // Disconnect - the HC might have closed the connection without us noticing and is asking for a reconnect
            final Connection connection = connectionManager.getConnection();
            StreamUtils.safeClose(connection);
            if(connection != null) {
                try {
                    // Wait for the connection to be closed
                    connection.awaitClosed();
                } catch (InterruptedException e) {
                    throw new InterruptedIOException();
                }
            }
        }

        boolean ok = false;
        final Connection connection = internalConnect();
        try {
            // Reconnect to the host-controller
            final ActiveOperation<Boolean, Void> result = channelHandler.executeRequest(new ServerReconnectRequest(), null);
            try {
                boolean inSync = result.getResult().get();
                ok = true;
                reconnectRunner = null;
                return inSync;
            } catch (ExecutionException e) {
                throw new IOException(e);
            } catch (InterruptedException e) {
                throw new InterruptedIOException();
            }
        } finally {
            if(!ok) {
                StreamUtils.safeClose(connection);
            }
        }
    }

    /**
     * Send the started notification
     */
    synchronized void started() {
        try {
            if(isConnected()) {
                channelHandler.executeRequest(new ServerStartedRequest(), null).getResult().await();
            }
        } catch (Exception e) {
            ServerLogger.AS_ROOT_LOGGER.debugf(e, "failed to send started notification");
        }
    }

    @Override
    public void connectionOpened(final Connection connection) throws IOException {
        final Channel channel = openChannel(connection, SERVER_CHANNEL_TYPE, configuration.getOptionMap());
        if(setChannel(channel)) {
            channel.receiveMessage(channelHandler.getReceiver());
            channel.addCloseHandler(channelHandler);
        } else {
            channel.closeAsync();
        }
    }

    @Override
    public void close() throws IOException {
        try {
            super.close();
            final ReconnectRunner reconnectTask = this.reconnectRunner;
            if (reconnectTask != null) {
                this.reconnectRunner = null;
                reconnectTask.cancel();
            }
        } finally {
            connectionManager.shutdown();
        }
    }

    /**
     * The server registration request.
     */
    private class ServerRegisterRequest extends AbstractManagementRequest<ModelNode, Void> {

        @Override
        public byte getOperationType() {
            return DomainServerProtocol.REGISTER_REQUEST;
        }

        @Override
        protected void sendRequest(final ActiveOperation.ResultHandler<ModelNode> resultHandler, final ManagementRequestContext<Void> context, final FlushableDataOutput output) throws IOException {
            output.writeUTF(serverProcessName);
            output.writeInt(initialOperationID);
        }

        @Override
        public void handleRequest(DataInput input, ActiveOperation.ResultHandler<ModelNode> resultHandler, ManagementRequestContext<Void> voidManagementRequestContext) throws IOException {
            final byte param = input.readByte();
            if(param == DomainServerProtocol.PARAM_OK) {
                final ModelNode operations = new ModelNode();
                operations.readExternal(input);
                resultHandler.done(operations);
            } else {
                resultHandler.failed(new IOException());
            }
        }

    }

    /**
     * The server reconnect request. Additionally to registering the server at the HC, the response will
     * contain whether this server is still in sync or needs to be restarted.
     */
    public class ServerReconnectRequest extends AbstractManagementRequest<Boolean, Void> {

        @Override
        public byte getOperationType() {
            return DomainServerProtocol.SERVER_RECONNECT_REQUEST;
        }

        @Override
        protected void sendRequest(final ActiveOperation.ResultHandler<Boolean> resultHandler, final ManagementRequestContext<Void> context, final FlushableDataOutput output) throws IOException {
            output.write(DomainServerProtocol.PARAM_SERVER_NAME);
            output.writeUTF(serverProcessName);
        }

        @Override
        public void handleRequest(final DataInput input, final ActiveOperation.ResultHandler<Boolean> resultHandler, final ManagementRequestContext<Void> context) throws IOException {
            final byte param = input.readByte();
            context.executeAsync(new ManagementRequestContext.AsyncTask<Void>() {
                @Override
                public void execute(ManagementRequestContext<Void> voidManagementRequestContext) throws Exception {
                    if(param == DomainServerProtocol.PARAM_OK) {
                        // Still in sync with the HC
                        resultHandler.done(Boolean.TRUE);
                    } else {
                        // Out of sync, set restart-required
                        resultHandler.done(Boolean.FALSE);
                    }
                }
            }, false);
        }

    }

    public class ServerStartedRequest extends AbstractManagementRequest<Void, Void> {

        private final String message = ""; // started / failed message

        @Override
        public byte getOperationType() {
            return DomainServerProtocol.SERVER_STARTED_REQUEST;
        }

        @Override
        protected void sendRequest(ActiveOperation.ResultHandler<Void> resultHandler, ManagementRequestContext<Void> voidManagementRequestContext, FlushableDataOutput output) throws IOException {
            output.write(DomainServerProtocol.PARAM_OK); // TODO handle server start failed message
            output.writeUTF(message);
            resultHandler.done(null);
        }

        @Override
        public void handleRequest(DataInput input, ActiveOperation.ResultHandler<Void> resultHandler, ManagementRequestContext<Void> voidManagementRequestContext) throws IOException {
            //
        }

    }

    private class ReconnectTask implements ProtocolConnectionManager.ConnectTask {

        @Override
        public Connection connect() throws IOException {
            // Reconnect with a potentially new configuration
            return ProtocolConnectionUtils.connectSync(configuration);
        }

        @Override
        public ProtocolConnectionManager.ConnectionOpenHandler getConnectionOpenedHandler() {
            return HostControllerConnection.this;
        }

        @Override
        public ProtocolConnectionManager.ConnectTask connectionClosed() {
            ServerLogger.AS_ROOT_LOGGER.debugf("Connection to Host Controller closed");
            return this;
        }

        @Override
        public void shutdown() {
            //
        }

    }

    interface ReconnectCallback {

        /**
         * Callback on reconnection.
         *
         * @param inSync    whether the server is still in sync with the host-controller
         */
        void reconnected(final boolean inSync);

    }

    class ReconnectRunner implements Runnable {

        private volatile Future<?> future;
        private volatile ReconnectCallback callback;

        @Override
        public synchronized void run() {
            final boolean outcome;
            try {
                outcome = doReConnect();
                callback.reconnected(outcome);
                reconnectRunner = null;
            } catch (Exception e) {
                try {
                    Thread.sleep(reconnectionDelay);
                } catch (InterruptedException i) {
                    Thread.currentThread().interrupt();
                }
                if (getState() == State.OPEN) {
                    ServerLogger.AS_ROOT_LOGGER.failedToConnectToHostController();
                    future = executorService.submit(this);
                }
            }
        }

        public void cancel() {
            if (future != null) {
                future.cancel(true);
            }
        }
    }

}
