/*
 * Copyright 2009 Mustard Grain, Inc., 2009-2010 LinkedIn, Inc.
 * 
 * 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 voldemort.server.niosocket;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;

import org.apache.commons.lang.mutable.MutableInt;
import org.apache.log4j.Level;

import voldemort.VoldemortException;
import voldemort.client.protocol.RequestFormatType;
import voldemort.server.protocol.RequestHandler;
import voldemort.server.protocol.RequestHandlerFactory;
import voldemort.server.protocol.StreamRequestHandler;
import voldemort.server.protocol.StreamRequestHandler.StreamRequestDirection;
import voldemort.server.protocol.StreamRequestHandler.StreamRequestHandlerState;
import voldemort.utils.ByteUtils;
import voldemort.utils.SelectorManagerWorker;

/**
 * AsyncRequestHandler manages a Selector, SocketChannel, and RequestHandler
 * implementation. At the point that the run method is invoked, the Selector
 * with which the (socket) Channel has been registered has notified us that the
 * socket has data to read or write.
 * <p/>
 * The bulk of the complexity in this class surrounds partial reads and writes,
 * as well as determining when all the data needed for the request has been
 * read.
 * 
 * 
 * @see voldemort.server.protocol.RequestHandler
 */

public class AsyncRequestHandler extends SelectorManagerWorker {

    private final RequestHandlerFactory requestHandlerFactory;

    private RequestHandler requestHandler;

    private StreamRequestHandler streamRequestHandler;

    private MutableInt serverConnectionCount;

    public AsyncRequestHandler(Selector selector,
                               SocketChannel socketChannel,
                               RequestHandlerFactory requestHandlerFactory,
                               int socketBufferSize,
                               MutableInt serverConnectionCount) {
        super(selector, socketChannel, socketBufferSize);
        this.requestHandlerFactory = requestHandlerFactory;
        this.serverConnectionCount = serverConnectionCount;
    }

    @Override
    protected void read(SelectionKey selectionKey) throws IOException {
        int count = 0;

        long startNs = -1;

        if(logger.isDebugEnabled())
            startNs = System.nanoTime();

        if((count = socketChannel.read(inputStream.getBuffer())) == -1)
            throw new EOFException("EOF for " + socketChannel.socket());

        if(logger.isTraceEnabled())
            traceInputBufferState("Read " + count + " bytes");

        if(count == 0)
            return;

        // Take note of the position after we read the bytes. We'll need it in
        // case of incomplete reads later on down the method.
        final int position = inputStream.getBuffer().position();

        // Flip the buffer, set our limit to the current position and then set
        // the position to 0 in preparation for reading in the RequestHandler.
        inputStream.getBuffer().flip();

        // We have to do this on the first request as we don't know the protocol
        // yet.
        if(requestHandler == null) {
            if(!initRequestHandler(selectionKey)) {
                return;
            }
        }

        if(streamRequestHandler != null) {
            // We're continuing an existing streaming request from our last pass
            // through. So handle it and return.
            handleStreamRequest(selectionKey);
            return;
        }

        if(!requestHandler.isCompleteRequest(inputStream.getBuffer())) {
            // Ouch - we're missing some data for a full request, so handle that
            // and return.
            handleIncompleteRequest(position);
            return;
        }

        // At this point we have the full request (and it's not streaming), so
        // rewind the buffer for reading and execute the request.
        inputStream.getBuffer().rewind();

        if(logger.isTraceEnabled())
            logger.trace("Starting execution for " + socketChannel.socket());

        DataInputStream dataInputStream = new DataInputStream(inputStream);
        DataOutputStream dataOutputStream = new DataOutputStream(outputStream);

        streamRequestHandler = requestHandler.handleRequest(dataInputStream,
							    dataOutputStream);

        if(logger.isDebugEnabled()) {
            logger.debug("AsyncRequestHandler:read finished request from "
                         + socketChannel.socket().getRemoteSocketAddress() + " handlerRef: "
                         + System.identityHashCode(dataInputStream) + " at time: "
                         + System.currentTimeMillis() + " elapsed time: "
                         + (System.nanoTime() - startNs) + " ns");
        }

        if(streamRequestHandler != null) {
            // In the case of a StreamRequestHandler, we handle that separately
            // (attempting to process multiple "segments").
            handleStreamRequest(selectionKey);
            return;
        }

        // At this point we've completed a full stand-alone request. So clear
        // our input buffer and prepare for outputting back to the client.
        if(logger.isTraceEnabled())
            logger.trace("Finished execution for " + socketChannel.socket());

        prepForWrite(selectionKey);
    }

    @Override
    protected void write(SelectionKey selectionKey) throws IOException {
        if(outputStream.getBuffer().hasRemaining()) {
            // If we have data, write what we can now...
            try {
                int count = socketChannel.write(outputStream.getBuffer());

                if(logger.isTraceEnabled())
                    logger.trace("Wrote " + count + " bytes, remaining: "
                                 + outputStream.getBuffer().remaining() + " for "
                                 + socketChannel.socket());
            } catch(IOException e) {
                if(streamRequestHandler != null) {
                    streamRequestHandler.close(new DataOutputStream(outputStream));
                    streamRequestHandler = null;
                }

                throw e;
            }

        } else {
            if(logger.isTraceEnabled())
                logger.trace("Wrote no bytes for " + socketChannel.socket());
        }

        // If there's more to write but we didn't write it, we'll take that to
        // mean that we're done here. We don't clear or reset anything. We leave
        // our buffer state where it is and try our luck next time.
        if(outputStream.getBuffer().hasRemaining())
            return;

        // If we don't have anything else to write, that means we're done with
        // the request! So clear the buffers (resizing if necessary).
        if(outputStream.getBuffer().capacity() >= resizeThreshold)
            outputStream.setBuffer(ByteBuffer.allocate(socketBufferSize));
        else
            outputStream.getBuffer().clear();

        if(streamRequestHandler != null
           && streamRequestHandler.getDirection() == StreamRequestDirection.WRITING) {
            // In the case of streaming writes, it's possible we can process
            // another segment of the stream. We process streaming writes this
            // way because there won't be any other notification for us to do
            // work as we won't be notified via reads.
            if(logger.isTraceEnabled())
                logger.trace("Request is streaming for " + socketChannel.socket());

            handleStreamRequest(selectionKey);
        } else {
            // If we're not streaming writes, signal the Selector that we're
            // ready to read the next request.
            selectionKey.interestOps(SelectionKey.OP_READ);
        }
    }

    private void handleStreamRequest(SelectionKey selectionKey) throws IOException {
        // You are not expected to understand this.
        DataInputStream dataInputStream = new DataInputStream(inputStream);
        DataOutputStream dataOutputStream = new DataOutputStream(outputStream);

        // We need to keep track of the last known starting index *before* we
        // attempt to service the next segment. This is needed in case of
        // partial reads so that we can revert back to this point.
        int preRequestPosition = inputStream.getBuffer().position();

        StreamRequestHandlerState state = handleStreamRequestInternal(selectionKey,
                                                                      dataInputStream,
                                                                      dataOutputStream);

        if(state == StreamRequestHandlerState.READING) {
            // We've read our request and handled one segment, but we aren't
            // ready to write anything just yet as we're streaming reads from
            // the client. So let's keep executing segments as much as we can
            // until we're no longer reading anything.
            do {
                preRequestPosition = inputStream.getBuffer().position();
                state = handleStreamRequestInternal(selectionKey, dataInputStream, dataOutputStream);
            } while(state == StreamRequestHandlerState.READING);
        } else if(state == StreamRequestHandlerState.WRITING) {
            // We've read our request and written one segment, but we're still
            // ready to stream writes to the client. So let's keep executing
            // segments as much as we can until we're there's nothing more to do
            // or until we blow past our buffer.
            do {
                state = handleStreamRequestInternal(selectionKey, dataInputStream, dataOutputStream);
            } while(state == StreamRequestHandlerState.WRITING && !outputStream.wasExpanded());

            if(state != StreamRequestHandlerState.COMPLETE) {
                // We've read our request and are ready to start streaming
                // writes to the client.
                prepForWrite(selectionKey);
            }
        }

        if(state == null) {
            // We got an error...
            return;
        }

        if(state == StreamRequestHandlerState.INCOMPLETE_READ) {
            // We need the data that's in there so far and aren't ready to write
            // anything out yet, so don't clear the input buffer or signal that
            // we're ready to write. But we do want to compact the buffer as we
            // don't want it to trigger an increase in the buffer if we don't
            // need to do so.

            // We need to do the following steps...
            //
            // a) ...figure out where we are in the buffer...
            int currentPosition = inputStream.getBuffer().position();

            // b) ...position ourselves at the start of the incomplete
            // "segment"...
            inputStream.getBuffer().position(preRequestPosition);

            // c) ...then copy the data starting from preRequestPosition's data
            // is at index 0...
            inputStream.getBuffer().compact();

            // d) ...and reset the position to be ready for the rest of the
            // reads and the limit to allow more data.
            handleIncompleteRequest(currentPosition - preRequestPosition);
        } else if(state == StreamRequestHandlerState.COMPLETE) {
            streamRequestHandler.close(dataOutputStream);
            streamRequestHandler = null;

            // Treat this as a normal request. Assume that all completed
            // requests want to write something back to the client.
            prepForWrite(selectionKey);
        }
    }

    private StreamRequestHandlerState handleStreamRequestInternal(SelectionKey selectionKey,
                                                                  DataInputStream dataInputStream,
                                                                  DataOutputStream dataOutputStream)
            throws IOException {
        StreamRequestHandlerState state = null;

        try {
            if(logger.isTraceEnabled())
                traceInputBufferState("Before streaming request handler");

            // this is the lowest level in the NioSocketServer stack at which we
            // still have a reference to the client IP address and port
            long startNs = -1;

            if(logger.isDebugEnabled())
                startNs = System.nanoTime();

            state = streamRequestHandler.handleRequest(dataInputStream, dataOutputStream);

            if(logger.isDebugEnabled()) {
                logger.debug("Handled request from "
                             + socketChannel.socket().getRemoteSocketAddress() + " handlerRef: "
                             + System.identityHashCode(dataInputStream) + " at time: "
                             + System.currentTimeMillis() + " elapsed time: "
                             + (System.nanoTime() - startNs) + " ns");
            }

            if(logger.isTraceEnabled())
                traceInputBufferState("After streaming request handler");
        } catch(Exception e) {
            if(logger.isEnabledFor(Level.WARN))
                logger.warn(e.getMessage(), e);

            VoldemortException error = e instanceof VoldemortException ? (VoldemortException) e
                                                                      : new VoldemortException(e);
            streamRequestHandler.handleError(dataOutputStream, error);
            streamRequestHandler.close(dataOutputStream);
            streamRequestHandler = null;

            prepForWrite(selectionKey);

            close();
        }

        return state;
    }

    /**
     * Returns true if the request should continue.
     * 
     * @return
     */

    private boolean initRequestHandler(SelectionKey selectionKey) {
        ByteBuffer inputBuffer = inputStream.getBuffer();
        int remaining = inputBuffer.remaining();

        // Don't have enough bytes to determine the protocol yet...
        if(remaining < 3)
            return true;

        byte[] protoBytes = { inputBuffer.get(0), inputBuffer.get(1), inputBuffer.get(2) };

        try {
            String proto = ByteUtils.getString(protoBytes, "UTF-8");
            RequestFormatType requestFormatType = RequestFormatType.fromCode(proto);
            requestHandler = requestHandlerFactory.getRequestHandler(requestFormatType);

            if(logger.isInfoEnabled())
                logger.info("Protocol negotiated for " + socketChannel.socket() + ": "
                            + requestFormatType.getDisplayName());

            // The protocol negotiation is the first request, so respond by
            // sticking the bytes in the output buffer, signaling the Selector,
            // and returning false to denote no further processing is needed.
            outputStream.getBuffer().put(ByteUtils.getBytes("ok", "UTF-8"));
            prepForWrite(selectionKey);

            return false;
        } catch(IllegalArgumentException e) {
            // okay we got some nonsense. For backwards compatibility,
            // assume this is an old client who does not know how to negotiate
            RequestFormatType requestFormatType = RequestFormatType.VOLDEMORT_V0;
            requestHandler = requestHandlerFactory.getRequestHandler(requestFormatType);

            if(logger.isInfoEnabled())
                logger.info("No protocol proposal given for " + socketChannel.socket()
                            + ", assuming " + requestFormatType.getDisplayName());

            return true;
        }
    }

    @Override
    public void close() {
        if(!isClosed.compareAndSet(false, true))
            return;

        serverConnectionCount.decrement();
        closeInternal();
    }
}
