/*
 * 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.arrow.flight;

import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST;
import static org.apache.arrow.flight.Location.forGrpcInsecure;
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;

import io.grpc.Channel;
import io.grpc.MethodDescriptor;
import io.grpc.ServerServiceDefinition;
import io.grpc.health.v1.HealthCheckRequest;
import io.grpc.health.v1.HealthCheckResponse;
import io.grpc.health.v1.HealthGrpc;
import io.grpc.netty.NettyChannelBuilder;
import io.grpc.netty.NettyServerBuilder;
import io.grpc.protobuf.services.HealthStatusManager;
import java.io.File;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import org.apache.arrow.flight.TestBasicOperation.Producer;
import org.apache.arrow.flight.auth.ServerAuthHandler;
import org.apache.arrow.flight.impl.FlightServiceGrpc;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.Test;

public class TestServerOptions {

  @Test
  public void builderConsumer() throws Exception {
    final AtomicBoolean consumerCalled = new AtomicBoolean();
    final Consumer<NettyServerBuilder> consumer = (builder) -> consumerCalled.set(true);

    try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        Producer producer = new Producer(a);
        FlightServer s =
            FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer)
                .transportHint("grpc.builderConsumer", consumer)
                .build()
                .start()) {
      assertTrue(consumerCalled.get());
    }
  }

  /**
   * Make sure that if Flight supplies a default executor to gRPC, then it is closed along with the
   * server.
   */
  @Test
  public void defaultExecutorClosed() throws Exception {
    final ExecutorService executor;
    try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        FlightServer server =
            FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer())
                .build()
                .start()) {
      assertNotNull(server.grpcExecutor);
      executor = server.grpcExecutor;
    }
    assertTrue(executor.isShutdown());
  }

  /** Make sure that if the user provides an executor to gRPC, then Flight does not close it. */
  @Test
  public void suppliedExecutorNotClosed() throws Exception {
    final ExecutorService executor = Executors.newSingleThreadExecutor();
    try {
      try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
          FlightServer server =
              FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), new NoOpFlightProducer())
                  .executor(executor)
                  .build()
                  .start()) {
        assertNull(server.grpcExecutor);
      }
      assertFalse(executor.isShutdown());
    } finally {
      executor.shutdown();
    }
  }

  @Test
  public void domainSocket() throws Exception {
    Assumptions.assumeTrue(
        FlightTestUtil.isNativeTransportAvailable(), "We have a native transport available");
    final File domainSocket = File.createTempFile("flight-unit-test-", ".sock");
    assertTrue(domainSocket.delete());
    // Domain socket paths have a platform-dependent limit. Set a conservative limit and skip the
    // test if the temporary
    // file name is too long. (We do not assume a particular platform-dependent temporary directory
    // path.)
    Assumptions.assumeTrue(
        domainSocket.getAbsolutePath().length() < 100, "The domain socket path is not too long");
    final Location location = Location.forGrpcDomainSocket(domainSocket.getAbsolutePath());
    try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        Producer producer = new Producer(a);
        FlightServer s = FlightServer.builder(a, location, producer).build().start(); ) {
      try (FlightClient c = FlightClient.builder(a, location).build()) {
        try (FlightStream stream = c.getStream(new Ticket(new byte[0]))) {
          VectorSchemaRoot root = stream.getRoot();
          IntVector iv = (IntVector) root.getVector("c1");
          int value = 0;
          while (stream.next()) {
            for (int i = 0; i < root.getRowCount(); i++) {
              assertEquals(value, iv.get(i));
              value++;
            }
          }
        }
      }
    }
  }

  @Test
  public void checkReflectionMetadata() {
    // This metadata is needed for gRPC reflection to work.
    final ExecutorService executorService = Executors.newSingleThreadExecutor();
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
      final FlightBindingService service =
          new FlightBindingService(
              allocator, new NoOpFlightProducer(), ServerAuthHandler.NO_OP, executorService);
      final ServerServiceDefinition definition = service.bindService();
      assertEquals(
          FlightServiceGrpc.getServiceDescriptor().getSchemaDescriptor(),
          definition.getServiceDescriptor().getSchemaDescriptor());

      final Map<String, MethodDescriptor<?, ?>> definedMethods = new HashMap<>();
      final Map<String, MethodDescriptor<?, ?>> serviceMethods = new HashMap<>();

      // Make sure that the reflection metadata object is identical across all the places where it's
      // accessible
      definition
          .getMethods()
          .forEach(
              method ->
                  definedMethods.put(
                      method.getMethodDescriptor().getFullMethodName(),
                      method.getMethodDescriptor()));
      definition
          .getServiceDescriptor()
          .getMethods()
          .forEach(method -> serviceMethods.put(method.getFullMethodName(), method));

      for (final MethodDescriptor<?, ?> descriptor :
          FlightServiceGrpc.getServiceDescriptor().getMethods()) {
        final String methodName = descriptor.getFullMethodName();
        assertTrue(
            definedMethods.containsKey(methodName),
            "Method is missing from ServerServiceDefinition: " + methodName);
        assertTrue(
            definedMethods.containsKey(methodName),
            "Method is missing from ServiceDescriptor: " + methodName);

        assertEquals(
            descriptor.getSchemaDescriptor(), definedMethods.get(methodName).getSchemaDescriptor());
        assertEquals(
            descriptor.getSchemaDescriptor(), serviceMethods.get(methodName).getSchemaDescriptor());
      }
    } finally {
      executorService.shutdown();
    }
  }

  /*
   * This is an extension of builderConsumer test.
   * Test that Flight interceptors don't break other registered services
   */
  @Test
  public void addHealthCheckService() throws Exception {
    final HealthStatusManager statusManager = new HealthStatusManager();
    final Consumer<NettyServerBuilder> consumer =
        (builder) -> {
          builder.addService(statusManager.getHealthService());
        };
    final Location location = forGrpcInsecure(LOCALHOST, 5555);
    try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        Producer producer = new Producer(a);
        FlightServer s =
            FlightServer.builder(a, location, producer)
                .transportHint("grpc.builderConsumer", consumer)
                .build()
                .start(); ) {
      Channel channel =
          NettyChannelBuilder.forAddress(location.toSocketAddress()).usePlaintext().build();
      HealthCheckResponse response =
          HealthGrpc.newBlockingStub(channel).check(HealthCheckRequest.getDefaultInstance());

      assertEquals(response.getStatus(), HealthCheckResponse.ServingStatus.SERVING);
    }
  }
}
