// Copyright (c) 2007-Present Pivotal Software, Inc.  All rights reserved.
//
// This software, the RabbitMQ Java client library, is triple-licensed under the
// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
// please see LICENSE-APACHE2.
//
// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
// either express or implied. See the LICENSE file for specific language governing
// rights and limitations of this software.
//
// If you have any questions regarding licensing, please contact us at
// info@rabbitmq.com.

package com.rabbitmq.client.impl.nio;

import com.rabbitmq.client.impl.Environment;
import com.rabbitmq.client.impl.Frame;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadFactory;

/**
 * Logic of the NIO loop.
 */
public class NioLoop implements Runnable {

    private static final Logger LOGGER = LoggerFactory.getLogger(NioLoop.class);

    private final NioLoopContext context;

    private final NioParams nioParams;

    public NioLoop(NioParams nioParams, NioLoopContext loopContext) {
        this.nioParams = nioParams;
        this.context = loopContext;
    }

    @Override
    public void run() {
        final SelectorHolder selectorState = context.readSelectorState;
        final Selector selector = selectorState.selector;
        final Set<SocketChannelRegistration> registrations = selectorState.registrations;

        final ByteBuffer buffer = context.readBuffer;

        final SelectorHolder writeSelectorState = context.writeSelectorState;
        final Selector writeSelector = writeSelectorState.selector;
        final Set<SocketChannelRegistration> writeRegistrations = writeSelectorState.registrations;

        // whether there have been write registrations in the previous loop
        // registrations are done after Selector.select(), to work on clean keys
        // thus, any write operation is performed in the next loop
        // we don't want to wait in the read Selector.select() if there are
        // pending writes
        boolean writeRegistered = false;

        try {
            while (true && !Thread.currentThread().isInterrupted()) {

                for (SelectionKey selectionKey : selector.keys()) {
                    SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) selectionKey.attachment();
                    if (state.getConnection() != null && state.getConnection().getHeartbeat() > 0) {
                        long now = System.currentTimeMillis();
                        if ((now - state.getLastActivity()) > state.getConnection().getHeartbeat() * 1000 * 2) {
                            try {
                                state.getConnection().handleHeartbeatFailure();
                            } catch (Exception e) {
                                LOGGER.warn("Error after heartbeat failure of connection {}", state.getConnection());
                            } finally {
                                selectionKey.cancel();
                            }
                        }
                    }
                }

                int select;
                if (!writeRegistered && registrations.isEmpty() && writeRegistrations.isEmpty()) {
                    // we can block, registrations will call Selector.wakeup()
                    select = selector.select(1000);
                    if (selector.keys().size() == 0) {
                        // we haven't been doing anything for a while, shutdown state
                        boolean clean = context.cleanUp();
                        if (clean) {
                            // we stop this thread
                            return;
                        }
                        // there may be incoming connections, keep going
                    }
                } else {
                    // we don't have to block, we need to select and clean cancelled keys before registration
                    select = selector.selectNow();
                }

                writeRegistered = false;

                // registrations should be done after select,
                // once the cancelled keys have been actually removed
                SocketChannelRegistration registration;
                Iterator<SocketChannelRegistration> registrationIterator = registrations.iterator();
                while (registrationIterator.hasNext()) {
                    registration = registrationIterator.next();
                    registrationIterator.remove();
                    int operations = registration.operations;
                    registration.state.getChannel().register(selector, operations, registration.state);
                }

                if (select > 0) {
                    Set<SelectionKey> readyKeys = selector.selectedKeys();
                    Iterator<SelectionKey> iterator = readyKeys.iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        iterator.remove();

                        if (!key.isValid()) {
                            continue;
                        }

                        if (key.isReadable()) {
                            final SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) key.attachment();

                            try {
                                if (!state.getChannel().isOpen()) {
                                    key.cancel();
                                    continue;
                                }
                                if(state.getConnection() == null) {
                                    // we're in AMQConnection#start, between the header sending and the FrameHandler#initialize
                                    // let's wait a bit more
                                    continue;
                                }

                                DataInputStream inputStream = state.inputStream;

                                state.prepareForReadSequence();

                                while (state.continueReading()) {
                                    Frame frame = Frame.readFrom(inputStream);

                                    try {
                                        boolean noProblem = state.getConnection().handleReadFrame(frame);
                                        if (noProblem && (!state.getConnection().isRunning() || state.getConnection().hasBrokerInitiatedShutdown())) {
                                            // looks like the frame was Close-Ok or Close
                                            dispatchShutdownToConnection(state);
                                            key.cancel();
                                            break;
                                        }
                                    } catch (Throwable ex) {
                                        // problem during frame processing, tell connection, and
                                        // we can stop for this channel
                                        handleIoError(state, ex);
                                        key.cancel();
                                        break;
                                    }
                                }

                                state.setLastActivity(System.currentTimeMillis());
                            } catch (final Exception e) {
                                LOGGER.warn("Error during reading frames", e);
                                handleIoError(state, e);
                                key.cancel();
                            } finally {
                                buffer.clear();
                            }
                        }
                    }
                }

                // write loop

                select = writeSelector.selectNow();

                // registrations should be done after select,
                // once the cancelled keys have been actually removed
                SocketChannelRegistration writeRegistration;
                Iterator<SocketChannelRegistration> writeRegistrationIterator = writeRegistrations.iterator();
                while (writeRegistrationIterator.hasNext()) {
                    writeRegistration = writeRegistrationIterator.next();
                    writeRegistrationIterator.remove();
                    int operations = writeRegistration.operations;
                    try {
                        if (writeRegistration.state.getChannel().isOpen()) {
                            writeRegistration.state.getChannel().register(writeSelector, operations, writeRegistration.state);
                            writeRegistered = true;
                        }
                    } catch (Exception e) {
                        // can happen if the channel has been closed since the operation has been enqueued
                        LOGGER.info("Error while registering socket channel for write: {}", e.getMessage());
                    }
                }

                if (select > 0) {
                    Set<SelectionKey> readyKeys = writeSelector.selectedKeys();
                    Iterator<SelectionKey> iterator = readyKeys.iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        iterator.remove();
                        SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) key.attachment();

                        if (!key.isValid()) {
                            continue;
                        }

                        if (key.isWritable()) {
                            boolean cancelKey = true;
                            try {
                                if (!state.getChannel().isOpen()) {
                                    key.cancel();
                                    continue;
                                }

                                state.prepareForWriteSequence();

                                int toBeWritten = state.getWriteQueue().size();
                                int written = 0;

                                DataOutputStream outputStream = state.outputStream;

                                WriteRequest request;
                                while (written <= toBeWritten && (request = state.getWriteQueue().poll()) != null) {
                                    request.handle(outputStream);
                                    written++;
                                }
                                outputStream.flush();
                                if (!state.getWriteQueue().isEmpty()) {
                                    cancelKey = true;
                                }
                            } catch (Exception e) {
                                handleIoError(state, e);
                            } finally {
                                state.endWriteSequence();
                                if (cancelKey) {
                                    key.cancel();
                                }
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            LOGGER.error("Error in NIO loop", e);
        }
    }

    protected void handleIoError(SocketChannelFrameHandlerState state, Throwable ex) {
        if (needToDispatchIoError(state)) {
            dispatchIoErrorToConnection(state, ex);
        } else {
            try {
                state.close();
            } catch (IOException e) {

            }
        }
    }

    protected boolean needToDispatchIoError(final SocketChannelFrameHandlerState state) {
        return state.getConnection().isOpen();
    }

    protected void dispatchIoErrorToConnection(final SocketChannelFrameHandlerState state, final Throwable ex) {
        // In case of recovery after the shutdown,
        // the new connection shouldn't be initialized in
        // the NIO thread, to avoid a deadlock.
        Runnable shutdown = () -> state.getConnection().handleIoError(ex);
        if (executorService() == null) {
            String name = "rabbitmq-connection-shutdown-" + state.getConnection();
            Thread shutdownThread = Environment.newThread(threadFactory(), shutdown, name);
            shutdownThread.start();
        } else {
            executorService().submit(shutdown);
        }
    }

    protected void dispatchShutdownToConnection(final SocketChannelFrameHandlerState state) {
        Runnable shutdown = new Runnable() {

            @Override
            public void run() {
                state.getConnection().doFinalShutdown();
            }
        };
        if (executorService() == null) {
            String name = "rabbitmq-connection-shutdown-" + state.getConnection();
            Thread shutdownThread = Environment.newThread(threadFactory(), shutdown, name);
            shutdownThread.start();
        } else {
            executorService().submit(shutdown);
        }
    }

    private ExecutorService executorService() {
        return nioParams.getNioExecutor();
    }

    private ThreadFactory threadFactory() {
        return nioParams.getThreadFactory();
    }
}
