/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.apache.catalina.tribes.transport.nio;
import java.io.IOException;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.DatagramChannel;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.nio.channels.WritableByteChannel;

import org.apache.catalina.tribes.ChannelMessage;
import org.apache.catalina.tribes.ChannelReceiver;
import org.apache.catalina.tribes.RemoteProcessException;
import org.apache.catalina.tribes.UniqueId;
import org.apache.catalina.tribes.io.BufferPool;
import org.apache.catalina.tribes.io.ChannelData;
import org.apache.catalina.tribes.io.ListenCallback;
import org.apache.catalina.tribes.io.ObjectReader;
import org.apache.catalina.tribes.transport.AbstractRxTask;
import org.apache.catalina.tribes.transport.Constants;
import org.apache.catalina.tribes.util.Logs;

/**
 * A worker thread class which can drain channels and echo-back the input. Each
 * instance is constructed with a reference to the owning thread pool object.
 * When started, the thread loops forever waiting to be awakened to service the
 * channel associated with a SelectionKey object. The worker is tasked by
 * calling its serviceChannel() method with a SelectionKey object. The
 * serviceChannel() method stores the key reference in the thread object then
 * calls notify() to wake it up. When the channel has been drained, the worker
 * thread returns itself to its parent pool.
 *
 * @author Filip Hanik
 *
 * @version $Id: NioReplicationTask.java 1142666 2011-07-04 13:54:52Z kkolinko $
 */
public class NioReplicationTask extends AbstractRxTask {

    private static final org.apache.juli.logging.Log log = org.apache.juli.logging.LogFactory.getLog( NioReplicationTask.class );

    private ByteBuffer buffer = null;
    private SelectionKey key;
    private int rxBufSize;
    private NioReceiver receiver;
    public NioReplicationTask (ListenCallback callback, NioReceiver receiver)
    {
        super(callback);
        this.receiver = receiver;
    }

    // loop forever waiting for work to do
    @Override
    public synchronized void run() {
        if ( buffer == null ) {
            int size = getRxBufSize();
            if (key.channel() instanceof DatagramChannel) {
                size = ChannelReceiver.MAX_UDP_SIZE;
            }
            if ( (getOptions() & OPTION_DIRECT_BUFFER) == OPTION_DIRECT_BUFFER) {
                buffer = ByteBuffer.allocateDirect(size);
            } else {
                buffer = ByteBuffer.allocate(size);
            }
        } else {
            buffer.clear();
        }
        if (key == null) {
            return; // just in case
        }
        if ( log.isTraceEnabled() )
            log.trace("Servicing key:"+key);

        try {
            ObjectReader reader = (ObjectReader)key.attachment();
            if ( reader == null ) {
                if ( log.isTraceEnabled() )
                    log.trace("No object reader, cancelling:"+key);
                cancelKey(key);
            } else {
                if ( log.isTraceEnabled() )
                    log.trace("Draining channel:"+key);

                drainChannel(key, reader);
            }
        } catch (Exception e) {
            //this is common, since the sockets on the other
            //end expire after a certain time.
            if ( e instanceof CancelledKeyException ) {
                //do nothing
            } else if ( e instanceof IOException ) {
                //dont spew out stack traces for IO exceptions unless debug is enabled.
                if (log.isDebugEnabled()) log.debug ("IOException in replication worker, unable to drain channel. Probable cause: Keep alive socket closed["+e.getMessage()+"].", e);
                else log.warn ("IOException in replication worker, unable to drain channel. Probable cause: Keep alive socket closed["+e.getMessage()+"].");
            } else if ( log.isErrorEnabled() ) {
                //this is a real error, log it.
                log.error("Exception caught in TcpReplicationThread.drainChannel.",e);
            }
            cancelKey(key);
        } finally {

        }
        key = null;
        // done, ready for more, return to pool
        getTaskPool().returnWorker (this);
    }

    /**
     * Called to initiate a unit of work by this worker thread
     * on the provided SelectionKey object.  This method is
     * synchronized, as is the run() method, so only one key
     * can be serviced at a given time.
     * Before waking the worker thread, and before returning
     * to the main selection loop, this key's interest set is
     * updated to remove OP_READ.  This will cause the selector
     * to ignore read-readiness for this channel while the
     * worker thread is servicing it.
     */
    public synchronized void serviceChannel (SelectionKey key) {
        if ( log.isTraceEnabled() ) log.trace("About to service key:"+key);
        ObjectReader reader = (ObjectReader)key.attachment();
        if ( reader != null ) reader.setLastAccess(System.currentTimeMillis());
        this.key = key;
        key.interestOps (key.interestOps() & (~SelectionKey.OP_READ));
        key.interestOps (key.interestOps() & (~SelectionKey.OP_WRITE));
    }

    /**
     * The actual code which drains the channel associated with
     * the given key.  This method assumes the key has been
     * modified prior to invocation to turn off selection
     * interest in OP_READ.  When this method completes it
     * re-enables OP_READ and calls wakeup() on the selector
     * so the selector will resume watching this channel.
     */
    protected void drainChannel (final SelectionKey key, ObjectReader reader) throws Exception {
        reader.setLastAccess(System.currentTimeMillis());
        reader.access();
        ReadableByteChannel channel = (ReadableByteChannel) key.channel();
        int count=-1;
        buffer.clear();         // make buffer empty
        SocketAddress saddr = null;

        if (channel instanceof SocketChannel) {
            // loop while data available, channel is non-blocking
            while ((count = channel.read (buffer)) > 0) {
                buffer.flip();      // make buffer readable
                if ( buffer.hasArray() )
                    reader.append(buffer.array(),0,count,false);
                else
                    reader.append(buffer,count,false);
                buffer.clear();     // make buffer empty
                //do we have at least one package?
                if ( reader.hasPackage() ) break;
            }
        } else if (channel instanceof DatagramChannel) {
            DatagramChannel dchannel = (DatagramChannel)channel;
            saddr = dchannel.receive(buffer);
            buffer.flip();      // make buffer readable
            if ( buffer.hasArray() )
                reader.append(buffer.array(),0,buffer.limit()-buffer.position(),false);
            else
                reader.append(buffer,buffer.limit()-buffer.position(),false);
            buffer.clear();     // make buffer empty
            //did we get a package
            count = reader.hasPackage()?1:-1;
        }

        int pkgcnt = reader.count();

        if (count < 0 && pkgcnt == 0 ) {
            //end of stream, and no more packages to process
            remoteEof(key);
            return;
        }

        ChannelMessage[] msgs = pkgcnt == 0? ChannelData.EMPTY_DATA_ARRAY : reader.execute();

        registerForRead(key,reader);//register to read new data, before we send it off to avoid dead locks

        for ( int i=0; i<msgs.length; i++ ) {
            /**
             * Use send ack here if you want to ack the request to the remote
             * server before completing the request
             * This is considered an asynchronized request
             */
            if (ChannelData.sendAckAsync(msgs[i].getOptions())) sendAck(key,(WritableByteChannel)channel,Constants.ACK_COMMAND,saddr);
            try {
                if ( Logs.MESSAGES.isTraceEnabled() ) {
                    try {
                        Logs.MESSAGES.trace("NioReplicationThread - Received msg:" + new UniqueId(msgs[i].getUniqueId()) + " at " + new java.sql.Timestamp(System.currentTimeMillis()));
                    }catch ( Throwable t ) {}
                }
                //process the message
                getCallback().messageDataReceived(msgs[i]);
                /**
                 * Use send ack here if you want the request to complete on this
                 * server before sending the ack to the remote server
                 * This is considered a synchronized request
                 */
                if (ChannelData.sendAckSync(msgs[i].getOptions())) sendAck(key,(WritableByteChannel)channel,Constants.ACK_COMMAND,saddr);
            }catch ( RemoteProcessException e ) {
                if ( log.isDebugEnabled() ) log.error("Processing of cluster message failed.",e);
                if (ChannelData.sendAckSync(msgs[i].getOptions())) sendAck(key,(WritableByteChannel)channel,Constants.FAIL_ACK_COMMAND,saddr);
            }catch ( Exception e ) {
                log.error("Processing of cluster message failed.",e);
                if (ChannelData.sendAckSync(msgs[i].getOptions())) sendAck(key,(WritableByteChannel)channel,Constants.FAIL_ACK_COMMAND,saddr);
            }
            if ( getUseBufferPool() ) {
                BufferPool.getBufferPool().returnBuffer(msgs[i].getMessage());
                msgs[i].setMessage(null);
            }
        }

        if (count < 0) {
            remoteEof(key);
            return;
        }
    }

    private void remoteEof(SelectionKey key) {
        // close channel on EOF, invalidates the key
        if ( log.isDebugEnabled() ) log.debug("Channel closed on the remote end, disconnecting");
        cancelKey(key);
    }

    protected void registerForRead(final SelectionKey key, ObjectReader reader) {
        if ( log.isTraceEnabled() )
            log.trace("Adding key for read event:"+key);
        reader.finish();
        //register our OP_READ interest
        Runnable r = new Runnable() {
            @Override
            public void run() {
                try {
                    if (key.isValid()) {
                        // cycle the selector so this key is active again
                        key.selector().wakeup();
                        // resume interest in OP_READ, OP_WRITE
                        int resumeOps = key.interestOps() | SelectionKey.OP_READ;
                        key.interestOps(resumeOps);
                        if ( log.isTraceEnabled() )
                            log.trace("Registering key for read:"+key);
                    }
                } catch (CancelledKeyException ckx ) {
                    NioReceiver.cancelledKey(key);
                    if ( log.isTraceEnabled() )
                        log.trace("CKX Cancelling key:"+key);

                } catch (Exception x) {
                    log.error("Error registering key for read:"+key,x);
                }
            }
        };
        receiver.addEvent(r);
    }

    private void cancelKey(final SelectionKey key) {
        if ( log.isTraceEnabled() )
            log.trace("Adding key for cancel event:"+key);

        ObjectReader reader = (ObjectReader)key.attachment();
        if ( reader != null ) {
            reader.setCancelled(true);
            reader.finish();
        }
        Runnable cx = new Runnable() {
            @Override
            public void run() {
                if ( log.isTraceEnabled() )
                    log.trace("Cancelling key:"+key);

                NioReceiver.cancelledKey(key);
            }
        };
        receiver.addEvent(cx);
    }


    /**
     * send a reply-acknowledgement (6,2,3), sends it doing a busy write, the ACK is so small
     * that it should always go to the buffer
     * @param key
     * @param channel
     */
    protected void sendAck(SelectionKey key, WritableByteChannel channel, byte[] command, SocketAddress udpaddr) {
        try {

            ByteBuffer buf = ByteBuffer.wrap(command);
            int total = 0;
            if (channel instanceof DatagramChannel) {
                DatagramChannel dchannel = (DatagramChannel)channel;
                //were using a shared channel, document says its thread safe
                //TODO check optimization, one channel per thread?
                while ( total < command.length ) {
                    total += dchannel.send(buf, udpaddr);
                }
            } else {
                while ( total < command.length ) {
                    total += channel.write(buf);
                }
            }
            if (log.isTraceEnabled()) {
                log.trace("ACK sent to " +
                        ( (channel instanceof SocketChannel) ?
                          ((SocketChannel)channel).socket().getInetAddress() :
                          ((DatagramChannel)channel).socket().getInetAddress()));
            }
        } catch ( java.io.IOException x ) {
            log.warn("Unable to send ACK back through channel, channel disconnected?: "+x.getMessage());
        }
    }

    public void setRxBufSize(int rxBufSize) {
        this.rxBufSize = rxBufSize;
    }

    public int getRxBufSize() {
        return rxBufSize;
    }
}
