/*
    * 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.maven.plugin.surefire.extensions;
  
  import javax.annotation.Nonnull;
  
  import java.io.Closeable;
  import java.io.IOException;
  import java.net.InetAddress;
  import java.net.InetSocketAddress;
  import java.net.SocketOption;
  import java.nio.Buffer;
  import java.nio.ByteBuffer;
  import java.nio.channels.AsynchronousServerSocketChannel;
  import java.nio.channels.AsynchronousSocketChannel;
  import java.nio.channels.ReadableByteChannel;
  import java.nio.channels.WritableByteChannel;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  
  import org.apache.maven.plugin.surefire.booterclient.output.NativeStdOutStreamConsumer;
  import org.apache.maven.surefire.api.event.Event;
  import org.apache.maven.surefire.api.fork.ForkNodeArguments;
  import org.apache.maven.surefire.extensions.CloseableDaemonThread;
  import org.apache.maven.surefire.extensions.CommandReader;
  import org.apache.maven.surefire.extensions.EventHandler;
  import org.apache.maven.surefire.extensions.ForkChannel;
  import org.apache.maven.surefire.extensions.util.CountDownLauncher;
  import org.apache.maven.surefire.extensions.util.CountdownCloseable;
  import org.apache.maven.surefire.extensions.util.LineConsumerThread;
  
  import static java.net.StandardSocketOptions.SO_KEEPALIVE;
  import static java.net.StandardSocketOptions.SO_REUSEADDR;
  import static java.net.StandardSocketOptions.TCP_NODELAY;
  import static java.nio.channels.AsynchronousChannelGroup.withThreadPool;
  import static java.nio.channels.AsynchronousServerSocketChannel.open;
  import static java.nio.charset.StandardCharsets.US_ASCII;
  import static org.apache.maven.surefire.api.util.internal.Channels.newBufferedChannel;
  import static org.apache.maven.surefire.api.util.internal.Channels.newChannel;
  import static org.apache.maven.surefire.api.util.internal.Channels.newInputStream;
  import static org.apache.maven.surefire.api.util.internal.Channels.newOutputStream;
  import static org.apache.maven.surefire.api.util.internal.DaemonThreadFactory.newDaemonThreadFactory;
  import static org.apache.maven.surefire.shared.lang3.StringUtils.isBlank;
  import static org.apache.maven.surefire.shared.lang3.StringUtils.isNotBlank;
  
  /**
   * The TCP/IP server accepting only one client connection. The forked JVM connects to the server using the
   * {@link #getForkNodeConnectionString() connection string}.
   * The main purpose of this class is to {@link #tryConnectToClient() conect with tthe client}, bind the
   * {@link #bindCommandReader(CommandReader, WritableByteChannel) command reader} to the internal socket's
   * {@link java.io.InputStream}, and bind the
   * {@link #bindEventHandler(EventHandler, CountdownCloseable, ReadableByteChannel) event handler} writing the event
   * objects to the {@link EventHandler event handler}.
   * <br>
   * The objects {@link WritableByteChannel} and {@link ReadableByteChannel} are forked process streams
   * (standard input and output). Both are ignored in this implementation but they are used in {@link LegacyForkChannel}.
   * <br>
   * The channel is closed after the forked JVM has finished normally or the shutdown hook is executed in the plugin.
   */
  final class SurefireForkChannel extends ForkChannel {
      private static final ExecutorService THREAD_POOL = Executors.newCachedThreadPool(newDaemonThreadFactory());
  
      private final AsynchronousServerSocketChannel server;
      private final String localHost;
      private final int localPort;
      private final String sessionId;
      private final Bindings bindings = new Bindings(2);
      private volatile Future<AsynchronousSocketChannel> session;
      private volatile LineConsumerThread out;
      private volatile CloseableDaemonThread commandReaderBindings;
      private volatile CloseableDaemonThread eventHandlerBindings;
      private volatile EventBindings eventBindings;
      private volatile CommandBindings commandBindings;
  
      SurefireForkChannel(@Nonnull ForkNodeArguments arguments) throws IOException {
          super(arguments);
          server = open(withThreadPool(THREAD_POOL));
          setTrueOptions(SO_REUSEADDR, TCP_NODELAY, SO_KEEPALIVE);
          InetAddress ip = InetAddress.getLoopbackAddress();
          server.bind(new InetSocketAddress(ip, 0), 1);
          InetSocketAddress localAddress = (InetSocketAddress) server.getLocalAddress();
         localHost = localAddress.getHostString();
         localPort = localAddress.getPort();
         sessionId = arguments.getSessionId();
     }
 
     @Override
     public void tryConnectToClient() {
         if (session != null) {
             throw new IllegalStateException("already accepted TCP client connection");
         }
         session = server.accept();
     }
 
     @Override
     public String getForkNodeConnectionString() {
         return "tcp://" + localHost + ":" + localPort + (isBlank(sessionId) ? "" : "?sessionId=" + sessionId);
     }
 
     @Override
     public int getCountdownCloseablePermits() {
         return 3;
     }
 
     @Override
     public void bindCommandReader(@Nonnull CommandReader commands, WritableByteChannel stdIn)
             throws IOException, InterruptedException {
         commandBindings = new CommandBindings(commands);
 
         bindings.countDown();
     }
 
     @Override
     public void bindEventHandler(
             @Nonnull EventHandler<Event> eventHandler,
             @Nonnull CountdownCloseable countdown,
             ReadableByteChannel stdOut)
             throws IOException, InterruptedException {
         ForkNodeArguments args = getArguments();
         out = new LineConsumerThread(
                 "fork-" + args.getForkChannelId() + "-out-thread",
                 stdOut,
                 new NativeStdOutStreamConsumer(args.getConsoleLock()),
                 countdown);
         out.start();
 
         eventBindings = new EventBindings(eventHandler, countdown);
 
         bindings.countDown();
     }
 
     @Override
     public void disable() {
         if (eventHandlerBindings != null) {
             eventHandlerBindings.disable();
         }
 
         if (commandReaderBindings != null) {
             commandReaderBindings.disable();
         }
     }
 
     @Override
     public void close() throws IOException {
         //noinspection unused,EmptyTryBlock,EmptyTryBlock
         try (Closeable c1 = getChannel();
                 Closeable c2 = server;
                 Closeable c3 = out) {
             // only close all channels
         } catch (InterruptedException e) {
             Throwable cause = e.getCause();
             throw cause instanceof IOException ? (IOException) cause : new IOException(cause);
         }
     }
 
     private void verifySessionId() throws InterruptedException, IOException {
         try {
             ByteBuffer buffer = ByteBuffer.allocate(sessionId.length());
             int read;
             do {
                 read = getChannel().read(buffer).get();
             } while (read != -1 && buffer.hasRemaining());
 
             if (read == -1) {
                 throw new IOException("Channel closed while verifying the client.");
             }
 
             ((Buffer) buffer).flip();
             String clientSessionId = new String(buffer.array(), US_ASCII);
             if (!clientSessionId.equals(sessionId)) {
                 throw new InvalidSessionIdException(clientSessionId, sessionId);
             }
         } catch (ExecutionException e) {
             Throwable cause = e.getCause();
             throw cause instanceof IOException ? (IOException) cause : new IOException(cause);
         }
     }
 
     @SafeVarargs
     private final void setTrueOptions(SocketOption<Boolean>... options) throws IOException {
         for (SocketOption<Boolean> option : options) {
             if (server.supportedOptions().contains(option)) {
                 server.setOption(option, true);
             }
         }
     }
 
     private class EventBindings {
         private final EventHandler<Event> eventHandler;
         private final CountdownCloseable countdown;
 
         private EventBindings(EventHandler<Event> eventHandler, CountdownCloseable countdown) {
             this.eventHandler = eventHandler;
             this.countdown = countdown;
         }
 
         void bindEventHandler(AsynchronousSocketChannel source) {
             ForkNodeArguments args = getArguments();
             String threadName = "fork-" + args.getForkChannelId() + "-event-thread";
             ReadableByteChannel channel = newBufferedChannel(newInputStream(source));
             eventHandlerBindings = new EventConsumerThread(threadName, channel, eventHandler, countdown, args);
             eventHandlerBindings.start();
         }
     }
 
     private class CommandBindings {
         private final CommandReader commands;
 
         private CommandBindings(CommandReader commands) {
             this.commands = commands;
         }
 
         void bindCommandSender(AsynchronousSocketChannel source) {
             // don't use newBufferedChannel here - may cause the command is not sent and the JVM hangs
             // only newChannel flushes the message
             // newBufferedChannel does not flush
             ForkNodeArguments args = getArguments();
             WritableByteChannel channel = newChannel(newOutputStream(source));
             String threadName = "commands-fork-" + args.getForkChannelId();
             commandReaderBindings = new StreamFeeder(threadName, channel, commands, args.getConsoleLogger());
             commandReaderBindings.start();
         }
     }
 
     private class Bindings extends CountDownLauncher {
         private Bindings(int count) {
             super(count);
         }
 
         @Override
         protected void job() throws IOException, InterruptedException {
             AsynchronousSocketChannel channel = getChannel();
             if (isNotBlank(sessionId)) {
                 verifySessionId();
             }
             eventBindings.bindEventHandler(channel);
             commandBindings.bindCommandSender(channel);
         }
     }
 
     private AsynchronousSocketChannel getChannel() throws InterruptedException, IOException {
         try {
             return session == null ? null : session.get();
         } catch (ExecutionException e) {
             Throwable cause = e.getCause();
             throw cause instanceof IOException ? (IOException) cause : new IOException(cause);
         }
     }
 }