/*
 * Copyright 2011 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.gradle.messaging.remote.internal;

import org.gradle.internal.CompositeStoppable;
import org.gradle.internal.id.IdGenerator;
import org.gradle.internal.concurrent.AsyncStoppable;
import org.gradle.internal.concurrent.ExecutorFactory;
import org.gradle.internal.concurrent.StoppableExecutor;
import org.gradle.messaging.dispatch.DiscardingFailureHandler;
import org.gradle.messaging.dispatch.Dispatch;
import org.gradle.messaging.dispatch.DispatchFailureHandler;
import org.gradle.messaging.remote.internal.protocol.EndOfStreamEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class MessageHub implements AsyncStoppable {
    private final Lock lock = new ReentrantLock();
    private final CompositeStoppable executors = CompositeStoppable.stoppable();
    private final CompositeStoppable connections = CompositeStoppable.stoppable();
    private final Collection<ProtocolStack<Message>> handlers = new ArrayList<ProtocolStack<Message>>();
    private final Collection<ProtocolStack<Message>> workers = new ArrayList<ProtocolStack<Message>>();
    private final Map<String, ProtocolStack<Message>> outgoingUnicasts = new HashMap<String, ProtocolStack<Message>>();
    private final Map<String, ProtocolStack<Message>> outgoingBroadcasts = new HashMap<String, ProtocolStack<Message>>();
    private final DispatchFailureHandler<Object> failureHandler;
    private final Router router;
    private final String displayName;
    private final String nodeName;
    private final ExecutorFactory executorFactory;
    private final IdGenerator<UUID> idGenerator;
    private final ClassLoader messagingClassLoader;
    private final StoppableExecutor incomingExecutor;

    public MessageHub(String displayName, String nodeName, ExecutorFactory executorFactory, IdGenerator<UUID> idGenerator, ClassLoader messagingClassLoader) {
        this.displayName = displayName;
        this.nodeName = nodeName;
        this.executorFactory = executorFactory;
        this.idGenerator = idGenerator;
        this.messagingClassLoader = messagingClassLoader;
        failureHandler = new DiscardingFailureHandler<Object>(LoggerFactory.getLogger(MessageHub.class));
        StoppableExecutor executor = executorFactory.create(displayName + " message router");
        executors.add(executor);
        router = new Router(executor, failureHandler);

        incomingExecutor = executorFactory.create(displayName + " worker");
        executors.add(incomingExecutor);
    }

    /**
     * Adds an incoming connection. Stops the connection when finished with it.
     */
    public void addConnection(Connection<Message> connection) {
        lock.lock();
        try {
            Connection<Message> wrapper = new EndOfStreamConnection(connection);
            AsyncConnectionAdapter<Message> asyncConnection = new AsyncConnectionAdapter<Message>(wrapper, failureHandler, executorFactory, new RemoteDisconnectProtocol());
            connections.add(asyncConnection);

            AsyncConnection<Message> incomingEndpoint = router.createRemoteConnection();
            incomingEndpoint.dispatchTo(new MethodInvocationMarshallingDispatch(asyncConnection));
            asyncConnection.dispatchTo(new MethodInvocationUnmarshallingDispatch(incomingEndpoint, messagingClassLoader));
        } finally {
            lock.unlock();
        }
    }

    public Dispatch<Object> addUnicastOutgoing(String channel) {
        lock.lock();
        try {
            ProtocolStack<Message> outgoing = outgoingUnicasts.get(channel);
            if (outgoing == null) {
                Protocol<Message> unicastSendProtocol = new UnicastSendProtocol();
                Protocol<Message> sendProtocol = new SendProtocol(idGenerator.generateId(), nodeName, channel);
                StoppableExecutor executor = executorFactory.create(displayName + " outgoing " + channel);
                executors.add(executor);
                outgoing = new ProtocolStack<Message>(executor, failureHandler, failureHandler, unicastSendProtocol, sendProtocol);
                outgoingUnicasts.put(channel, outgoing);

                AsyncConnection<Message> outgoingEndpoint = router.createLocalConnection();
                outgoing.getBottom().dispatchTo(outgoingEndpoint);
                outgoingEndpoint.dispatchTo(outgoing.getBottom());
            }
            return new OutgoingMultiplex(channel, outgoing.getTop());
        } finally {
            lock.unlock();
        }
    }

    public Dispatch<Object> addMulticastOutgoing(String channel) {
        lock.lock();
        try {
            ProtocolStack<Message> outgoing = outgoingBroadcasts.get(channel);
            if (outgoing == null) {
                Protocol<Message> broadcastProtocol = new BroadcastSendProtocol();
                Protocol<Message> sendProtocol = new SendProtocol(idGenerator.generateId(), nodeName, channel);
                StoppableExecutor executor = executorFactory.create(displayName + " outgoing broadcast " + channel);
                executors.add(executor);
                outgoing = new ProtocolStack<Message>(executor, failureHandler, failureHandler, broadcastProtocol, sendProtocol);
                outgoingBroadcasts.put(channel, outgoing);

                AsyncConnection<Message> outgoingEndpoint = router.createLocalConnection();
                outgoing.getBottom().dispatchTo(outgoingEndpoint);
                outgoingEndpoint.dispatchTo(outgoing.getBottom());
            }
            return new OutgoingMultiplex(channel, outgoing.getTop());
        } finally {
            lock.unlock();
        }
    }

    public void addIncoming(final String channel, final Dispatch<Object> dispatch) {
        lock.lock();
        try {
            final UUID id = idGenerator.generateId();
            Protocol<Message> workerProtocol = new WorkerProtocol(dispatch);
            Protocol<Message> receiveProtocol = new ReceiveProtocol(id, nodeName, channel);

            ProtocolStack<Message> workerStack = new ProtocolStack<Message>(incomingExecutor, failureHandler, failureHandler, workerProtocol);
            workers.add(workerStack);
            ProtocolStack<Message> stack = new ProtocolStack<Message>(incomingExecutor, failureHandler, failureHandler, new BufferingProtocol(200), receiveProtocol);
            handlers.add(stack);

            workerStack.getBottom().dispatchTo(stack.getTop());
            stack.getTop().dispatchTo(workerStack.getBottom());

            AsyncConnection<Message> incomingEndpoint = router.createLocalConnection();
            stack.getBottom().dispatchTo(incomingEndpoint);
            incomingEndpoint.dispatchTo(stack.getBottom());
        } finally {
            lock.unlock();
        }
    }

    public void requestStop() {
        lock.lock();
        try {
            for (ProtocolStack<Message> stack : outgoingUnicasts.values()) {
                stack.requestStop();
            }
            for (ProtocolStack<Message> stack : outgoingBroadcasts.values()) {
                stack.requestStop();
            }
            for (ProtocolStack<?> worker : workers) {
                worker.requestStop();
            }
        } finally {
            lock.unlock();
        }
    }

    public void stop() {
        requestStop();

        CompositeStoppable stoppable = new CompositeStoppable();
        lock.lock();
        try {
            stoppable.add(outgoingUnicasts.values());
            stoppable.add(outgoingBroadcasts.values());
            stoppable.add(workers);
            stoppable.add(handlers);
            stoppable.add(connections);
            stoppable.add(router);
            stoppable.add(executors);
        } finally {
            outgoingUnicasts.clear();
            outgoingBroadcasts.clear();
            workers.clear();
            handlers.clear();
            lock.unlock();
        }

        stoppable.stop();
    }

    private static class EndOfStreamConnection extends DelegatingConnection<Message> {
        private static final Logger LOGGER = LoggerFactory.getLogger(EndOfStreamConnection.class);
        boolean incomingFinished;

        private EndOfStreamConnection(Connection<Message> connection) {
            super(connection);
        }

        @Override
        public Message receive() {
            if (incomingFinished) {
                return null;
            }
            Message result;
            try {
                result = super.receive();
            } catch (Throwable e) {
                LOGGER.error("Could not receive message from connection. Discarding connection.", e);
                result = null;
            }
            if (result instanceof EndOfStreamEvent) {
                incomingFinished = true;
            } else if (result == null) {
                incomingFinished = true;
                result = new EndOfStreamEvent();
            }
            return result;
        }
    }
}
