/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch 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.elasticsearch.nio;

import org.elasticsearch.common.CheckedFunction;
import org.elasticsearch.common.util.BigArrays;
import org.elasticsearch.test.ESTestCase;
import org.junit.Before;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SocketChannel;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyInt;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class BytesChannelContextTests extends ESTestCase {

    private CheckedFunction<InboundChannelBuffer, Integer, IOException> readConsumer;
    private NioSocketChannel channel;
    private SocketChannel rawChannel;
    private BytesChannelContext context;
    private InboundChannelBuffer channelBuffer;
    private NioSelector selector;
    private BiConsumer<Void, Exception> listener;
    private int messageLength;

    @Before
    @SuppressWarnings("unchecked")
    public void init() {
        readConsumer = mock(CheckedFunction.class);

        messageLength = randomInt(96) + 20;
        selector = mock(NioSelector.class);
        listener = mock(BiConsumer.class);
        channel = mock(NioSocketChannel.class);
        rawChannel = mock(SocketChannel.class);
        channelBuffer = InboundChannelBuffer.allocatingInstance();
        TestReadWriteHandler handler = new TestReadWriteHandler(readConsumer);
        when(channel.getRawChannel()).thenReturn(rawChannel);
        context = new BytesChannelContext(channel, selector, mock(Consumer.class), handler, channelBuffer);

        when(selector.isOnCurrentThread()).thenReturn(true);
    }

    public void testSuccessfulRead() throws IOException {
        byte[] bytes = createMessage(messageLength);

        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenAnswer(invocationOnMock -> {
            ByteBuffer[] buffers = (ByteBuffer[]) invocationOnMock.getArguments()[0];
            buffers[0].put(bytes);
            return bytes.length;
        });

        when(readConsumer.apply(channelBuffer)).thenReturn(messageLength, 0);

        assertEquals(messageLength, context.read());

        assertEquals(0, channelBuffer.getIndex());
        assertEquals(BigArrays.BYTE_PAGE_SIZE - bytes.length, channelBuffer.getCapacity());
        verify(readConsumer, times(1)).apply(channelBuffer);
    }

    public void testMultipleReadsConsumed() throws IOException {
        byte[] bytes = createMessage(messageLength * 2);

        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenAnswer(invocationOnMock -> {
            ByteBuffer[] buffers = (ByteBuffer[]) invocationOnMock.getArguments()[0];
            buffers[0].put(bytes);
            return bytes.length;
        });

        when(readConsumer.apply(channelBuffer)).thenReturn(messageLength, messageLength, 0);

        assertEquals(bytes.length, context.read());

        assertEquals(0, channelBuffer.getIndex());
        assertEquals(BigArrays.BYTE_PAGE_SIZE - bytes.length, channelBuffer.getCapacity());
        verify(readConsumer, times(2)).apply(channelBuffer);
    }

    public void testPartialRead() throws IOException {
        byte[] bytes = createMessage(messageLength);

        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenAnswer(invocationOnMock -> {
            ByteBuffer[] buffers = (ByteBuffer[]) invocationOnMock.getArguments()[0];
            buffers[0].put(bytes);
            return bytes.length;
        });


        when(readConsumer.apply(channelBuffer)).thenReturn(0);

        assertEquals(messageLength, context.read());

        assertEquals(bytes.length, channelBuffer.getIndex());
        verify(readConsumer, times(1)).apply(channelBuffer);

        when(readConsumer.apply(channelBuffer)).thenReturn(messageLength * 2, 0);

        assertEquals(messageLength, context.read());

        assertEquals(0, channelBuffer.getIndex());
        assertEquals(BigArrays.BYTE_PAGE_SIZE - (bytes.length * 2), channelBuffer.getCapacity());
        verify(readConsumer, times(2)).apply(channelBuffer);
    }

    public void testReadThrowsIOException() throws IOException {
        IOException ioException = new IOException();
        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenThrow(ioException);

        IOException ex = expectThrows(IOException.class, () -> context.read());
        assertSame(ioException, ex);
    }

    public void testReadThrowsIOExceptionMeansReadyForClose() throws IOException {
        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenThrow(new IOException());

        assertFalse(context.selectorShouldClose());
        expectThrows(IOException.class, () -> context.read());
        assertTrue(context.selectorShouldClose());
    }

    public void testReadLessThanZeroMeansReadyForClose() throws IOException {
        when(rawChannel.read(any(ByteBuffer[].class), anyInt(), anyInt())).thenReturn(-1L);

        assertEquals(0, context.read());

        assertTrue(context.selectorShouldClose());
    }

    @SuppressWarnings("varargs")
    public void testQueuedWriteIsFlushedInFlushCall() throws Exception {
        assertFalse(context.readyForFlush());

        ByteBuffer[] buffers = {ByteBuffer.allocate(10)};

        FlushReadyWrite flushOperation = mock(FlushReadyWrite.class);
        context.queueWriteOperation(flushOperation);

        assertTrue(context.readyForFlush());

        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
        when(flushOperation.isFullyFlushed()).thenReturn(true);
        when(flushOperation.getListener()).thenReturn(listener);
        context.flushChannel();

        verify(rawChannel).write(buffers, 0, buffers.length);
        verify(selector).executeListener(listener, null);
        assertFalse(context.readyForFlush());
    }

    public void testPartialFlush() throws IOException {
        assertFalse(context.readyForFlush());
        FlushReadyWrite flushOperation = mock(FlushReadyWrite.class);
        context.queueWriteOperation(flushOperation);
        assertTrue(context.readyForFlush());

        when(flushOperation.isFullyFlushed()).thenReturn(false);
        when(flushOperation.getBuffersToWrite()).thenReturn(new ByteBuffer[0]);
        context.flushChannel();

        verify(listener, times(0)).accept(null, null);
        assertTrue(context.readyForFlush());
    }

    @SuppressWarnings("unchecked")
    public void testMultipleWritesPartialFlushes() throws IOException {
        assertFalse(context.readyForFlush());

        BiConsumer<Void, Exception> listener2 = mock(BiConsumer.class);
        FlushReadyWrite flushOperation1 = mock(FlushReadyWrite.class);
        FlushReadyWrite flushOperation2 = mock(FlushReadyWrite.class);
        when(flushOperation1.getBuffersToWrite()).thenReturn(new ByteBuffer[0]);
        when(flushOperation2.getBuffersToWrite()).thenReturn(new ByteBuffer[0]);
        when(flushOperation1.getListener()).thenReturn(listener);
        when(flushOperation2.getListener()).thenReturn(listener2);

        context.queueWriteOperation(flushOperation1);
        context.queueWriteOperation(flushOperation2);

        assertTrue(context.readyForFlush());

        when(flushOperation1.isFullyFlushed()).thenReturn(true);
        when(flushOperation2.isFullyFlushed()).thenReturn(false);
        context.flushChannel();

        verify(selector).executeListener(listener, null);
        verify(listener2, times(0)).accept(null, null);
        assertTrue(context.readyForFlush());

        when(flushOperation2.isFullyFlushed()).thenReturn(true);

        context.flushChannel();

        verify(selector).executeListener(listener2, null);
        assertFalse(context.readyForFlush());
    }

    public void testWhenIOExceptionThrownListenerIsCalled() throws IOException {
        assertFalse(context.readyForFlush());

        ByteBuffer[] buffers = {ByteBuffer.allocate(10)};
        FlushReadyWrite flushOperation = mock(FlushReadyWrite.class);
        context.queueWriteOperation(flushOperation);

        assertTrue(context.readyForFlush());

        IOException exception = new IOException();
        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
        when(rawChannel.write(buffers, 0, buffers.length)).thenThrow(exception);
        when(flushOperation.getListener()).thenReturn(listener);
        expectThrows(IOException.class, () -> context.flushChannel());

        verify(selector).executeFailedListener(listener, exception);
        assertFalse(context.readyForFlush());
    }

    public void testWriteIOExceptionMeansChannelReadyToClose() throws IOException {
        ByteBuffer[] buffers = {ByteBuffer.allocate(10)};
        FlushReadyWrite flushOperation = mock(FlushReadyWrite.class);
        context.queueWriteOperation(flushOperation);

        IOException exception = new IOException();
        when(flushOperation.getBuffersToWrite()).thenReturn(buffers);
        when(rawChannel.write(buffers, 0, buffers.length)).thenThrow(exception);

        assertFalse(context.selectorShouldClose());
        expectThrows(IOException.class, () -> context.flushChannel());
        assertTrue(context.selectorShouldClose());
    }

    public void testInitiateCloseSchedulesCloseWithSelector() {
        context.closeChannel();
        verify(selector).queueChannelClose(channel);
    }

    private static byte[] createMessage(int length) {
        byte[] bytes = new byte[length];
        for (int i = 0; i < length; ++i) {
            bytes[i] = randomByte();
        }
        return bytes;
    }

    private static class TestReadWriteHandler extends BytesWriteHandler {

        private final CheckedFunction<InboundChannelBuffer, Integer, IOException> fn;

        private TestReadWriteHandler(CheckedFunction<InboundChannelBuffer, Integer, IOException> fn) {
            this.fn = fn;
        }

        @Override
        public int consumeReads(InboundChannelBuffer channelBuffer) throws IOException {
            return fn.apply(channelBuffer);
        }
    }
}
