/*
      Copyright 2010 Ramon Servadei
   
      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 mirrormap.nio;
  
  import java.io.IOException;
  import java.net.InetSocketAddress;
  import java.nio.ByteBuffer;
  import java.nio.channels.SelectableChannel;
  import java.nio.channels.SelectionKey;
  import java.nio.channels.Selector;
  import java.nio.channels.SocketChannel;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Iterator;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.concurrent.Executor;
  import java.util.logging.Level;
  import java.util.logging.Logger;
  
  import mirrormap.Utils;
  import mirrormap.lifecycle.AbstractLifeCycle;
  
  /**
   * Handles the socket processing for a TCP end point.
   * 
   * @author Ramon Servadei
   */
  public final class TcpConnectionEndPoint extends AbstractLifeCycle implements
      IConnectionEndPoint, ISelectionKeyTask
  {
      /**
       * Bit shifts to access the byte ordinal that the index represents. E.g.
       * bitShiftForByteOrdinal[2] has a value of 8 which is the number of bits to
       * right shift an integral to access the 2nd byte.
       */
      private static final long[] bitShiftForByteOrdinal = new long[9];
      static
      {
          int i = 0;
          // 0 doesn't count
          bitShiftForByteOrdinal[i] = 0;
          bitShiftForByteOrdinal[++i] = 8 * (i - 1); // 0
          bitShiftForByteOrdinal[++i] = 8 * (i - 1); // 8
          bitShiftForByteOrdinal[++i] = 8 * (i - 1); // 16
          bitShiftForByteOrdinal[++i] = 8 * (i - 1);
          bitShiftForByteOrdinal[++i] = 8 * (i - 1);
          bitShiftForByteOrdinal[++i] = 8 * (i - 1);
          bitShiftForByteOrdinal[++i] = 8 * (i - 1);
          bitShiftForByteOrdinal[++i] = 8 * (i - 1);
      }
  
      private final static Logger LOGGER =
          Logger.getLogger(TcpConnectionEndPoint.class.getName());
  
      /** The receive buffer, 65535 bytes */
      private final static byte[] recvBuffer = new byte[65535];
  
      /** The pending messages to send */
      final List<ByteBuffer> messages = new ArrayList<ByteBuffer>();
  
      /**
       * The selection key associated with the {@link SelectableChannel} for this
       * end point
       */
      private SelectionKey selectionKey;
  
      /** Reference required during {@link #doComponentDestroy()} */
      private final SelectorTasks selectorTasks;
  
      /** Handles data received from the remote end point */
      private final IConnectionEndPointReceiver receiver;
  
      /** The socket channel to the remote {@link IConnectionEndPoint} */
      private SocketChannel socketChannel;
  
      /**
       * Flag to indicate the socket direction, by default connections are
       * outbound
       */
      private boolean outbound = true;
  
      /** The TCP port for the socket connection of the other end point */
      private final int remotePort;
  
      /** The IP address or resolvable host name of the other end point */
     private final String remoteHost;
 
     /** Handles the socket writing */
     private Executor writer;
 
     /**
      * Standard constructor for a TCP end point.
      * 
      * @param receiver
      *            the object that will handle data received from the remote end
      *            point
      * @param selectorTasks
      *            used to register and unregister this end point against the
      *            {@link SelectionKey} of the underlying
      *            {@link SelectableChannel} backing this end point
      * @param writer
      *            handles writing data, optional. If <code>null</code> then the
      *            {@link Thread} running the {@link SelectorTasks#process()}
      *            method will do the writing
      * @param host
      *            the IP address or resolvable host name of the other end point
      * @param port
      *            the TCP port for the socket connection of the other end point
      * @throws IllegalStateException
      *             if the socket could not be constructed
      */
     public TcpConnectionEndPoint(IConnectionEndPointReceiver receiver,
         SelectorTasks selectorTasks, Executor writer, String host, int port)
     {
         this.remoteHost = host;
         this.remotePort = port;
         this.receiver = receiver;
         this.selectorTasks = selectorTasks;
         this.writer = writer;
     }
 
     /**
      * Add the data to the pending messages list
      * 
      * @param data
      *            the data to add to the pending messages
      */
     public void send(byte[] data)
     {
         if (!isActive())
         {
             if (LOGGER.isLoggable(Level.FINE))
             {
                 LOGGER.fine("inactive");
             }
             return;
         }
         synchronized (this)
         {
             byte[] actual = encode(data);
 
             final ByteBuffer wrapped = ByteBuffer.wrap(actual);
             messages.add(wrapped);
             // we're interested in writing now
             this.selectionKey.interestOps(this.selectionKey.interestOps()
                 | SelectionKey.OP_WRITE);
             this.selectionKey.selector().wakeup();
         }
     }
 
     public IConnectionEndPointReceiver getReceiver()
     {
         return this.receiver;
     }
 
     @Override
     protected void doStart()
     {
         try
         {
             if (isOutbound())
             {
                 final SocketChannel localSocketChannel = SocketChannel.open();
                 localSocketChannel.configureBlocking(false);
                 // as the channel is in non-blocking mode, we need some extra
                 // work to connect
                 if (!localSocketChannel.connect(new InetSocketAddress(
                     this.remoteHost, this.remotePort)))
                 {
                     // the connection has not completed, so wait on the selector
                     // until we can complete the connection
                     final Selector selector = Selector.open();
                     final SelectionKey key =
                         localSocketChannel.register(selector,
                             SelectionKey.OP_CONNECT);
                     selector.select();
                     localSocketChannel.finishConnect();
                     key.cancel();
                     selector.close();
                 }
                 if (LOGGER.isLoggable(Level.FINE))
                 {
                     LOGGER.fine("(->) Connected to " + this);
                 }
                 setSocketChannel(localSocketChannel);
             }
             else
             {
                 if (LOGGER.isLoggable(Level.FINE))
                 {
                     LOGGER.fine("(<-) Connected to " + this);
                 }
             }
         }
         catch (Exception e)
         {
             this.socketChannel = null;
             final String warning =
                 "Could not create socket to " + this.remoteHost + ":"
                     + this.remotePort + ". Calling destroy().";
             destroy();
             throw new RuntimeException(warning, e);
         }
         this.selectorTasks.register(SelectionKey.OP_READ, getSocketChannel(),
             this);
     }
 
     public boolean isConnected()
     {
         return this.socketChannel != null
             && this.socketChannel.socket() != null
             && this.socketChannel.socket().isConnected();
     }
 
     /**
      * Set the direction of the socket connection.
      * 
      * @param outbound
      *            <code>true</code> for an outbound connecting socket (connects
      *            to a server socket), <code>false</code> for an inbound
      *            connecting one
      */
     void setOutbound(boolean outbound)
     {
         this.outbound = outbound;
     }
 
     /**
      * Indicates the connecting direction.
      * 
      * @return <code>true</code> if this socket connects to a server socket,
      *         <code>false</code> if it was an inbound connection from another
      *         end point
      */
     public boolean isOutbound()
     {
         return outbound;
     }
 
     /**
      * Destroy the internal TCP I/O components.
      */
     @Override
     protected void doDestroy()
     {
         if (this.receiver != null)
         {
             this.receiver.destroy();
         }
         if (this.socketChannel != null)
         {
             try
             {
                 if (this.socketChannel.socket() != null)
                 {
                     this.socketChannel.socket().getOutputStream().flush();
                     this.socketChannel.socket().shutdownInput();
                     this.socketChannel.socket().shutdownOutput();
                     this.socketChannel.socket().close();
                 }
                 this.socketChannel.close();
             }
             catch (IOException e)
             {
                 if (LOGGER.isLoggable(Level.FINE))
                 {
                     LOGGER.log(Level.FINE, "Could not close "
                         + Utils.safeToString(this), e);
                 }
             }
         }
         try
         {
             if (this.selectionKey != null
                 && this.selectionKey.channel() != null)
             {
                 this.selectorTasks.unregister(this.selectionKey);
                 this.selectionKey.channel().close();
             }
         }
         catch (Exception e)
         {
             Utils.logException(LOGGER, "Could not close channel for "
                 + Utils.safeToString(this), e);
         }
     }
 
     protected SocketChannel getSocketChannel()
     {
         return this.socketChannel;
     }
 
     void setSocketChannel(SocketChannel socketChannel)
     {
         this.socketChannel = socketChannel;
     }
 
     @Override
     public String toString()
     {
         return Utils.string(this, this.remoteHost + ":" + this.remotePort);
     }
 
     /**
      * Add a 4 byte prefix to the data that includes the length of the data as a
      * big-endian integer.
      * 
      * @param data
      *            the data
      * @return an array equal in length to <code>4 + data.length</code> with a
      *         copy of data starting from element 4
      */
     private byte[] encode(byte[] data)
     {
         // add the integer length of the data, always the first 4 bytes
         byte[] actual = new byte[data.length + 4];
         System.arraycopy(data, 0, actual, 4, data.length);
         int startFrom = 0;
         for (int i = 4; i > 0; i--)
         {
             actual[startFrom++] =
                 (byte) (data.length >>> bitShiftForByteOrdinal[i]);
         }
         return actual;
     }
 
     /**
      * Split the data into its sub-messages. Messages have a 4 byte prefix that
      * is the message length in bytes, excluding the 4 bytes, as a big-endian
      * integer. This allows multiple messages to be received as a single TCP
      * frame.
      * 
      * @param message
      *            the full TCP frame data
      * @return a list of the individual messages contained in the message
      *         argument
      */
     private List<byte[]> decode(byte[] message)
     {
         // the TCP message could be a concatenation of multiple ones
         // the format is: <4 bytes len><subdata><4 bytes len><subdata>
         List<byte[]> values = new LinkedList<byte[]>();
         int lenPtr = 0;
         int len = 0;
         while (lenPtr < message.length)
         {
             len = readInteger(message, lenPtr, 4);
             // copy the sub data into the new array
             byte[] subData = new byte[len];
             lenPtr += 4;
             System.arraycopy(message, lenPtr, subData, 0, len);
             lenPtr += len;
             values.add(subData);
         }
         return values;
     }
 
     /**
      * Reconstruct an integer from its byte form in the byte[].
      * 
      * @param bytes
      *            the byte[] holding the data in byte form
      * @param start
      *            the position in bytes where the data starts
      * @param numberOfBytes
      *            the number of bytes that form the data to read
      * @return an integer
      */
     private int readInteger(final byte[] bytes, final int start,
         int numberOfBytes)
     {
         long output = 0;
         final int end = numberOfBytes + start;
         // identify if this is a -ve number, has ramifications for building up
         // the output
         boolean negative = false;
         if ((bytes[start] & 0x80) == 0x80)
         {
             output = -1;
             negative = true;
         }
         for (int i = start; i < end; i++)
         {
             // cast to a char as it is unsigned, and also AND with 00FF to
             // correctly promote whilst keeping the byte unsigned
             final char dataByte = (char) (bytes[i] & 0x00FF);
             if (negative)
             {
                 // clear out this byte position with 0, prior to bitwise ORing
                 // the byte
                 output ^=
                     ((long) 0xff << bitShiftForByteOrdinal[numberOfBytes]);
             }
             // bitwise OR the byte after shifting it by the correct number of
             // bits
             output |=
                 ((long) dataByte << bitShiftForByteOrdinal[numberOfBytes--]);
         }
         return (int) output;
     }
 
     public void run()
     {
         if (isActive())
         {
             if (this.selectionKey.isWritable())
             {
                 write();
             }
             if (this.selectionKey.isReadable())
             {
                 read();
             }
         }
     }
 
     /**
      * Reads data from the socket
      */
     private void read()
     {
         try
         {
             final ByteBuffer recv = ByteBuffer.wrap(recvBuffer);
             final SocketChannel socketChannel =
                 ((SocketChannel) this.selectionKey.channel());
             final int size = socketChannel.read(recv);
             if (size == -1)
             {
                 /*
                  * end of stream reached (see
                  * ReadableByteChannel.read(ByteBuffer dst)
                  */
                 throw new IOException("end-of-stream");
             }
             final byte[] bufferData = new byte[recv.position()];
             System.arraycopy(recv.array(), 0, bufferData, 0, recv.position());
             // find any sub-messages
             final List<byte[]> decoded = decode(bufferData);
             for (byte[] byteData : decoded)
             {
                 if (LOGGER.isLoggable(Level.FINEST))
                 {
                     LOGGER.finest(this + " received '"
                         + Arrays.toString(byteData) + "'");
                 }
                 try
                 {
                     this.receiver.receive(byteData);
                 }
                 catch (Exception e)
                 {
                     Utils.logException(LOGGER, "Could not handle data", e);
                 }
             }
         }
         catch (IOException e)
         {
             Utils.logException(LOGGER, "Destroying " + Utils.safeToString(this)
                 + " because data could not be read from "
                 + Utils.safeToString(this.selectionKey.channel()), e);
             destroy();
         }
     }
 
     /**
      * Writes any data in the message buffer to the socket
      */
     private void write()
     {
         final List<ByteBuffer> copy =
             new ArrayList<ByteBuffer>(this.messages.size());
         synchronized (this)
         {
             copy.addAll(this.messages);
             this.messages.clear();
         }
         Runnable task = new Runnable()
         {
             public void run()
             {
                 for (Iterator<ByteBuffer> iterator = copy.iterator(); iterator.hasNext();)
                 {
                     final ByteBuffer data = iterator.next();
                     try
                     {
                         if (LOGGER.isLoggable(Level.FINEST))
                         {
                             LOGGER.finest(TcpConnectionEndPoint.this
                                 + " sending '" + Arrays.toString(data.array())
                                 + "'");
                         }
                         ((SocketChannel) TcpConnectionEndPoint.this.selectionKey.channel()).write(data);
                     }
                     catch (IOException e)
                     {
                         Utils.logException(LOGGER,
                             Utils.safeToString(TcpConnectionEndPoint.this)
                                 + " could not write data '"
                                 + new String(data.array()) + "'", e);
                     }
                 }
             }
         };
         if (this.writer == null)
         {
             task.run();
         }
         else
         {
             this.writer.execute(task);
         }
         synchronized (this)
         {
             if (this.messages.size() == 0)
             {
                 // nothing more to write, so remove write interest
                 this.selectionKey.interestOps(this.selectionKey.interestOps()
                     ^ SelectionKey.OP_WRITE);
             }
         }
     }
 
     public void setSelectionKey(SelectionKey selectionKey)
     {
         this.selectionKey = selectionKey;
     }
 
 }