/*
    * Copyright 2016 The Netty Project
    *
    * The Netty Project 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:
    *
    *   https://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 io.netty.testsuite.transport.socket;
  
  import io.netty.bootstrap.Bootstrap;
  import io.netty.bootstrap.ServerBootstrap;
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.CompositeByteBuf;
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelFuture;
  import io.netty.channel.ChannelFutureListener;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.channel.ChannelInboundHandlerAdapter;
  import io.netty.channel.ChannelInitializer;
  import io.netty.channel.ChannelOption;
  import io.netty.channel.socket.SocketChannel;
  import io.netty.util.concurrent.ImmediateEventExecutor;
  import io.netty.util.concurrent.Promise;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.TestInfo;
  import org.junit.jupiter.api.Timeout;
  
  import java.io.ByteArrayOutputStream;
  import java.net.InetSocketAddress;
  import java.net.SocketAddress;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.Semaphore;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Function;
  
  import static io.netty.buffer.ByteBufUtil.writeAscii;
  import static io.netty.buffer.UnpooledByteBufAllocator.DEFAULT;
  import static io.netty.util.CharsetUtil.US_ASCII;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  public class SocketConnectTest extends AbstractSocketTest {
  
      @Test
      @Timeout(value = 30000, unit = TimeUnit.MILLISECONDS)
      public void testCloseTwice(TestInfo testInfo) throws Throwable {
          run(testInfo, new Runner<ServerBootstrap, Bootstrap>() {
              @Override
              public void run(ServerBootstrap serverBootstrap, Bootstrap bootstrap) throws Throwable {
                  testCloseTwice(serverBootstrap, bootstrap);
              }
          });
      }
  
      public void testCloseTwice(ServerBootstrap sb, Bootstrap cb) throws Throwable {
          Channel serverChannel = null;
          Channel clientChannel = null;
          try {
              serverChannel = sb.childHandler(new ChannelInboundHandlerAdapter()).bind().syncUninterruptibly().channel();
              final BlockingQueue<ChannelFuture> futures = new LinkedBlockingQueue<>();
              clientChannel = cb.handler(new ChannelInboundHandlerAdapter() {
                          @Override
                          public void userEventTriggered(ChannelHandlerContext ctx, Object evt)  {
                              futures.add(ctx.close());
                          }
                      })
                      .connect(serverChannel.localAddress()).syncUninterruptibly().channel();
              clientChannel.pipeline().fireUserEventTriggered("test");
              clientChannel.close().syncUninterruptibly();
              futures.take().sync();
              clientChannel = null;
  
              serverChannel.close().syncUninterruptibly();
              serverChannel.close().syncUninterruptibly();
              serverChannel = null;
          } finally {
              if (clientChannel != null) {
                  clientChannel.close().syncUninterruptibly();
              }
              if (serverChannel != null) {
                  serverChannel.close().syncUninterruptibly();
              }
          }
      }
  
      @Test
      @Timeout(value = 30000, unit = TimeUnit.MILLISECONDS)
     public void testLocalAddressAfterConnect(TestInfo testInfo) throws Throwable {
         run(testInfo, new Runner<ServerBootstrap, Bootstrap>() {
             @Override
             public void run(ServerBootstrap serverBootstrap, Bootstrap bootstrap) throws Throwable {
                 testLocalAddressAfterConnect(serverBootstrap, bootstrap);
             }
         });
     }
 
     public void testLocalAddressAfterConnect(ServerBootstrap sb, Bootstrap cb) throws Throwable {
         Channel serverChannel = null;
         Channel clientChannel = null;
         try {
             final Promise<InetSocketAddress> localAddressPromise = ImmediateEventExecutor.INSTANCE.newPromise();
             serverChannel = sb.childHandler(new ChannelInboundHandlerAdapter() {
                         @Override
                         public void channelActive(ChannelHandlerContext ctx) throws Exception {
                             localAddressPromise.setSuccess((InetSocketAddress) ctx.channel().localAddress());
                         }
                     }).bind().syncUninterruptibly().channel();
 
             clientChannel = cb.handler(new ChannelInboundHandlerAdapter()).register().syncUninterruptibly().channel();
 
             assertNull(clientChannel.localAddress());
             assertNull(clientChannel.remoteAddress());
 
             clientChannel.connect(serverChannel.localAddress()).syncUninterruptibly().channel();
             assertLocalAddress((InetSocketAddress) clientChannel.localAddress());
             assertNotNull(clientChannel.remoteAddress());
 
             assertLocalAddress(localAddressPromise.get());
         } finally {
             if (clientChannel != null) {
                 clientChannel.close().syncUninterruptibly();
             }
             if (serverChannel != null) {
                 serverChannel.close().syncUninterruptibly();
             }
         }
     }
 
     @Test
     @Timeout(value = 3000, unit = TimeUnit.MILLISECONDS)
     public void testChannelEventsFiredWhenClosedDirectly(TestInfo testInfo) throws Throwable {
         run(testInfo, new Runner<ServerBootstrap, Bootstrap>() {
             @Override
             public void run(ServerBootstrap serverBootstrap, Bootstrap bootstrap) throws Throwable {
                 testChannelEventsFiredWhenClosedDirectly(serverBootstrap, bootstrap);
             }
         });
     }
 
     public void testChannelEventsFiredWhenClosedDirectly(ServerBootstrap sb, Bootstrap cb) throws Throwable {
         final BlockingQueue<Integer> events = new LinkedBlockingQueue<Integer>();
 
         Channel sc = null;
         Channel cc = null;
         try {
             sb.childHandler(new ChannelInboundHandlerAdapter());
             sc = sb.bind().syncUninterruptibly().channel();
 
             cb.handler(new ChannelInboundHandlerAdapter() {
                 @Override
                 public void channelActive(ChannelHandlerContext ctx) throws Exception {
                     events.add(0);
                 }
 
                 @Override
                 public void channelInactive(ChannelHandlerContext ctx) throws Exception {
                     events.add(1);
                 }
             });
             // Connect and directly close again.
             cc = cb.connect(sc.localAddress()).addListener(ChannelFutureListener.CLOSE).
                     syncUninterruptibly().channel();
             assertEquals(0, events.take().intValue());
             assertEquals(1, events.take().intValue());
         } finally {
             if (cc != null) {
                 cc.close();
             }
             if (sc != null) {
                 sc.close();
             }
         }
     }
 
     @Test
     @Timeout(value = 3000, unit = TimeUnit.MILLISECONDS)
     public void testWriteWithFastOpenBeforeConnect(TestInfo testInfo) throws Throwable {
         run(testInfo, new Runner<ServerBootstrap, Bootstrap>() {
             @Override
             public void run(ServerBootstrap serverBootstrap, Bootstrap bootstrap) throws Throwable {
                 testWriteWithFastOpenBeforeConnect(serverBootstrap, bootstrap);
             }
         });
     }
 
     @Test
     @Timeout(value = 3000, unit = TimeUnit.MILLISECONDS)
     public void testWriteWithFastOpenBeforeConnectDirectBufferReaderIndex(TestInfo testInfo) throws Throwable {
         run(testInfo, (serverBootstrap, bootstrap) ->
                 testWriteWithFastOpenBeforeConnect(serverBootstrap, bootstrap,
                         channel -> directBufferWithReaderIndex(
                                 channel, "BAD-PREFIX-", "[fastopen-reader-index]"),
                         "[fastopen-reader-index]"));
     }
 
     @Test
     @Timeout(value = 3000, unit = TimeUnit.MILLISECONDS)
     public void testWriteWithFastOpenBeforeConnectCompositeBuffer(TestInfo testInfo) throws Throwable {
         run(testInfo, (serverBootstrap, bootstrap) ->
                 testWriteWithFastOpenBeforeConnect(serverBootstrap, bootstrap,
                         channel -> compositeBuffer(channel, "[fastopen-", "composite-", "data]"),
                         "[fastopen-composite-data]"));
     }
 
     @Test
     @Timeout(value = 3000, unit = TimeUnit.MILLISECONDS)
     public void testWriteWithFastOpenBeforeConnectCompositeBufferReaderIndex(TestInfo testInfo) throws Throwable {
         run(testInfo, (serverBootstrap, bootstrap) ->
                 testWriteWithFastOpenBeforeConnect(serverBootstrap, bootstrap,
                         channel -> compositeBufferWithReaderIndex(
                                 channel, "BAD-PREFIX-", "[fastopen-", "composite-reader-", "index]"),
                         "[fastopen-composite-reader-index]"));
     }
 
     public void testWriteWithFastOpenBeforeConnect(ServerBootstrap sb, Bootstrap cb) throws Throwable {
         testWriteWithFastOpenBeforeConnect(sb, cb, channel -> writeAscii(DEFAULT, "[fastopen]"), "[fastopen]");
     }
 
     protected void testWriteWithFastOpenBeforeConnect(ServerBootstrap sb, Bootstrap cb,
                                                       Function<Channel, ByteBuf> initialDataFactory,
                                                       String expectedInitialData) throws Throwable {
         enableTcpFastOpen(sb, cb);
         sb.childOption(ChannelOption.AUTO_READ, true);
         cb.option(ChannelOption.AUTO_READ, true);
 
         sb.childHandler(new ChannelInitializer<SocketChannel>() {
             @Override
             protected void initChannel(SocketChannel ch) throws Exception {
                 ch.pipeline().addLast(new EchoServerHandler());
             }
         });
 
         Channel sc = sb.bind().sync().channel();
         connectAndVerifyDataTransfer(cb, sc, initialDataFactory, expectedInitialData);
         connectAndVerifyDataTransfer(cb, sc, initialDataFactory, expectedInitialData);
     }
 
     private static void connectAndVerifyDataTransfer(Bootstrap cb, Channel sc,
                                                      Function<Channel, ByteBuf> initialDataFactory,
                                                      String expectedInitialData)
             throws InterruptedException {
         BufferingClientHandler handler = new BufferingClientHandler();
         cb.handler(handler);
         ChannelFuture register = cb.register();
         Channel channel = register.sync().channel();
         ChannelFuture write = channel.write(initialDataFactory.apply(channel));
         SocketAddress remoteAddress = sc.localAddress();
         ChannelFuture connectFuture = channel.connect(remoteAddress);
         Channel cc = connectFuture.sync().channel();
         cc.writeAndFlush(writeAscii(DEFAULT, "[normal data]")).sync();
         write.sync();
         String expectedString = expectedInitialData + "[normal data]";
         String result = handler.collectBuffer(expectedString.getBytes(US_ASCII).length);
         cc.disconnect().sync();
         assertEquals(expectedString, result);
     }
 
     private static ByteBuf directBufferWithReaderIndex(Channel channel, String prefix, String data) {
         ByteBuf buffer = directBuffer(channel, prefix + data);
         buffer.readerIndex(prefix.getBytes(US_ASCII).length);
         return buffer;
     }
 
     private static ByteBuf compositeBufferWithReaderIndex(Channel channel, String prefix, String... parts) {
         CompositeByteBuf buffer = channel.alloc().compositeDirectBuffer(parts.length + 1);
         buffer.addComponent(true, directBuffer(channel, prefix));
         for (String part : parts) {
             buffer.addComponent(true, directBuffer(channel, part));
         }
         buffer.readerIndex(prefix.getBytes(US_ASCII).length);
         return buffer;
     }
 
     private static ByteBuf compositeBuffer(Channel channel, String... parts) {
         CompositeByteBuf buffer = channel.alloc().compositeDirectBuffer(parts.length);
         for (String part : parts) {
             buffer.addComponent(true, directBuffer(channel, part));
         }
         return buffer;
     }
 
     private static ByteBuf directBuffer(Channel channel, String data) {
         byte[] bytes = data.getBytes(US_ASCII);
         ByteBuf buffer = channel.alloc().directBuffer(bytes.length);
         buffer.writeBytes(bytes);
         return buffer;
     }
 
     protected void enableTcpFastOpen(ServerBootstrap sb, Bootstrap cb) {
         // TFO is an almost-pure optimisation and should not change any observable behaviour in our tests.
         sb.option(ChannelOption.TCP_FASTOPEN, 5);
         cb.option(ChannelOption.TCP_FASTOPEN_CONNECT, true);
     }
 
     private static void assertLocalAddress(InetSocketAddress address) {
         assertTrue(address.getPort() > 0);
         assertFalse(address.getAddress().isAnyLocalAddress());
     }
 
     private static class BufferingClientHandler extends ChannelInboundHandlerAdapter {
         private final Semaphore semaphore = new Semaphore(0);
         private final ByteArrayOutputStream streamBuffer = new ByteArrayOutputStream();
 
         @Override
         public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
             if (msg instanceof ByteBuf) {
                 ByteBuf buf = (ByteBuf) msg;
                 int readableBytes = buf.readableBytes();
                 buf.readBytes(streamBuffer, readableBytes);
                 semaphore.release(readableBytes);
                 buf.release();
             } else {
                 throw new IllegalArgumentException("Unexpected message type: " + msg);
             }
         }
 
         String collectBuffer(int expectedBytes) throws InterruptedException {
             semaphore.acquire(expectedBytes);
             byte[] bytes = streamBuffer.toByteArray();
             streamBuffer.reset();
             return new String(bytes, US_ASCII);
         }
     }
 
     private static final class EchoServerHandler extends ChannelInboundHandlerAdapter {
         @Override
         public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
             if (msg instanceof ByteBuf) {
                 ByteBuf buffer = ctx.alloc().buffer();
                 ByteBuf buf = (ByteBuf) msg;
                 buffer.writeBytes(buf);
                 buf.release();
                 ctx.channel().writeAndFlush(buffer);
             } else {
                 throw new IllegalArgumentException("Unexpected message type: " + msg);
             }
         }
     }
 }