/*
 * 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 java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/**
 * A basic test of client middleware using a simplified OpenTracing-like example.
 */
@RunWith(JUnit4.class)
public class TestClientMiddleware {

  /**
   * Test that a client middleware can fail a call before it starts by throwing a {@link FlightRuntimeException}.
   */
  @Test
  public void clientMiddleware_failCallBeforeSending() {
    test(new NoOpFlightProducer(), null, Collections.singletonList(new CallRejector.Factory()),
        (allocator, client) -> {
          FlightTestUtil.assertCode(FlightStatusCode.UNAVAILABLE, client::listActions);
        });
  }

  /**
   * Test an OpenTracing-like scenario where client and server middleware work together to propagate a request ID
   * without explicit intervention from the service implementation.
   */
  @Test
  public void middleware_propagateHeader() {
    final Context context = new Context("span id");
    test(new NoOpFlightProducer(),
        new TestServerMiddleware.ServerMiddlewarePair<>(
            FlightServerMiddleware.Key.of("test"), new ServerSpanInjector.Factory()),
        Collections.singletonList(new ClientSpanInjector.Factory(context)),
        (allocator, client) -> {
          FlightTestUtil.assertCode(FlightStatusCode.UNIMPLEMENTED, () -> client.listActions().forEach(actionType -> {
          }));
        });
    Assert.assertEquals(context.outgoingSpanId, context.incomingSpanId);
    Assert.assertNotNull(context.finalStatus);
    Assert.assertEquals(FlightStatusCode.UNIMPLEMENTED, context.finalStatus.code());
  }

  /** Ensure both server and client can send and receive multi-valued headers (both binary and text values). */
  @Test
  public void testMultiValuedHeaders() {
    final MultiHeaderClientMiddlewareFactory clientFactory = new MultiHeaderClientMiddlewareFactory();
    test(new NoOpFlightProducer(),
        new TestServerMiddleware.ServerMiddlewarePair<>(
            FlightServerMiddleware.Key.of("test"), new MultiHeaderServerMiddlewareFactory()),
        Collections.singletonList(clientFactory),
        (allocator, client) -> {
          FlightTestUtil.assertCode(FlightStatusCode.UNIMPLEMENTED, () -> client.listActions().forEach(actionType -> {
          }));
        });
    // The server echoes the headers we send back to us, so ensure all the ones we sent are present with the correct
    // values in the correct order.
    for (final Map.Entry<String, List<byte[]>> entry : EXPECTED_BINARY_HEADERS.entrySet()) {
      // Compare header values entry-by-entry because byte arrays don't compare via equals
      final List<byte[]> receivedValues = clientFactory.lastBinaryHeaders.get(entry.getKey());
      Assert.assertNotNull("Missing for header: " + entry.getKey(), receivedValues);
      Assert.assertEquals(
          "Missing or wrong value for header: " + entry.getKey(),
          entry.getValue().size(), receivedValues.size());
      for (int i = 0; i < entry.getValue().size(); i++) {
        Assert.assertArrayEquals(entry.getValue().get(i), receivedValues.get(i));
      }
    }
    for (final Map.Entry<String, List<String>> entry : EXPECTED_TEXT_HEADERS.entrySet()) {
      Assert.assertEquals(
          "Missing or wrong value for header: " + entry.getKey(),
          entry.getValue(), clientFactory.lastTextHeaders.get(entry.getKey()));
    }
  }

  private static <T extends FlightServerMiddleware> void test(FlightProducer producer,
      TestServerMiddleware.ServerMiddlewarePair<T> serverMiddleware,
      List<FlightClientMiddleware.Factory> clientMiddleware,
      BiConsumer<BufferAllocator, FlightClient> body) {
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
      final FlightServer server = FlightTestUtil
          .getStartedServer(location -> {
            final FlightServer.Builder builder = FlightServer.builder(allocator, location, producer);
            if (serverMiddleware != null) {
              builder.middleware(serverMiddleware.key, serverMiddleware.factory);
            }
            return builder.build();
          });
      FlightClient.Builder builder = FlightClient.builder(allocator, server.getLocation());
      clientMiddleware.forEach(builder::intercept);
      try (final FlightServer ignored = server;
          final FlightClient client = builder.build()
      ) {
        body.accept(allocator, client);
      }
    } catch (InterruptedException | IOException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * A server middleware component that reads a request ID from incoming headers and sends the request ID back on
   * outgoing headers.
   */
  static class ServerSpanInjector implements FlightServerMiddleware {

    private final String spanId;

    public ServerSpanInjector(String spanId) {
      this.spanId = spanId;
    }

    @Override
    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
      outgoingHeaders.insert("x-span", spanId);
    }

    @Override
    public void onCallCompleted(CallStatus status) {

    }

    @Override
    public void onCallErrored(Throwable err) {

    }

    static class Factory implements FlightServerMiddleware.Factory<ServerSpanInjector> {

      @Override
      public ServerSpanInjector onCallStarted(CallInfo info, CallHeaders incomingHeaders, RequestContext context) {
        return new ServerSpanInjector(incomingHeaders.get("x-span"));
      }
    }
  }

  /**
   * A client middleware component that, given a mock OpenTracing-like "request context", sends the request ID in the
   * context on outgoing headers and reads it from incoming headers.
   */
  static class ClientSpanInjector implements FlightClientMiddleware {

    private final Context context;

    public ClientSpanInjector(Context context) {
      this.context = context;
    }

    @Override
    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
      outgoingHeaders.insert("x-span", context.outgoingSpanId);
    }

    @Override
    public void onHeadersReceived(CallHeaders incomingHeaders) {
      context.incomingSpanId = incomingHeaders.get("x-span");
    }

    @Override
    public void onCallCompleted(CallStatus status) {
      context.finalStatus = status;
    }

    static class Factory implements FlightClientMiddleware.Factory {

      private final Context context;

      Factory(Context context) {
        this.context = context;
      }

      @Override
      public FlightClientMiddleware onCallStarted(CallInfo info) {
        return new ClientSpanInjector(context);
      }
    }
  }

  /**
   * A mock OpenTracing-like "request context".
   */
  static class Context {

    final String outgoingSpanId;
    String incomingSpanId;
    CallStatus finalStatus;

    Context(String spanId) {
      this.outgoingSpanId = spanId;
    }
  }

  /**
   * A client middleware that fails outgoing calls.
   */
  static class CallRejector implements FlightClientMiddleware {

    @Override
    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
    }

    @Override
    public void onHeadersReceived(CallHeaders incomingHeaders) {
    }

    @Override
    public void onCallCompleted(CallStatus status) {
    }

    static class Factory implements FlightClientMiddleware.Factory {

      @Override
      public FlightClientMiddleware onCallStarted(CallInfo info) {
        throw CallStatus.UNAVAILABLE.withDescription("Rejecting call.").toRuntimeException();
      }
    }
  }

  // Used to test that middleware can send and receive multi-valued text and binary headers.
  static final Map<String, List<byte[]>> EXPECTED_BINARY_HEADERS = new HashMap<String, List<byte[]>>() {{
      put("x-binary-bin", Arrays.asList(new byte[] {0}, new byte[]{1}));
    }};
  static final Map<String, List<String>> EXPECTED_TEXT_HEADERS = new HashMap<String, List<String>>() {{
      put("x-text", Arrays.asList("foo", "bar"));
    }};

  static class MultiHeaderServerMiddlewareFactory implements
      FlightServerMiddleware.Factory<MultiHeaderServerMiddleware> {
    @Override
    public MultiHeaderServerMiddleware onCallStarted(CallInfo info, CallHeaders incomingHeaders,
        RequestContext context) {
      // Echo the headers back to the client. Copy values out of CallHeaders since the underlying gRPC metadata
      // object isn't safe to use after this function returns.
      Map<String, List<byte[]>> binaryHeaders = new HashMap<>();
      Map<String, List<String>> textHeaders = new HashMap<>();
      for (final String key : incomingHeaders.keys()) {
        if (key.endsWith("-bin")) {
          binaryHeaders.compute(key, (ignored, values) -> {
            if (values == null) {
              values = new ArrayList<>();
            }
            incomingHeaders.getAllByte(key).forEach(values::add);
            return values;
          });
        } else {
          textHeaders.compute(key, (ignored, values) -> {
            if (values == null) {
              values = new ArrayList<>();
            }
            incomingHeaders.getAll(key).forEach(values::add);
            return values;
          });
        }
      }
      return new MultiHeaderServerMiddleware(binaryHeaders, textHeaders);
    }
  }

  static class MultiHeaderServerMiddleware implements FlightServerMiddleware {
    private final Map<String, List<byte[]>> binaryHeaders;
    private final Map<String, List<String>> textHeaders;

    MultiHeaderServerMiddleware(Map<String, List<byte[]>> binaryHeaders, Map<String, List<String>> textHeaders) {
      this.binaryHeaders = binaryHeaders;
      this.textHeaders = textHeaders;
    }

    @Override
    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
      binaryHeaders.forEach((key, values) -> values.forEach(value -> outgoingHeaders.insert(key, value)));
      textHeaders.forEach((key, values) -> values.forEach(value -> outgoingHeaders.insert(key, value)));
    }

    @Override
    public void onCallCompleted(CallStatus status) {}

    @Override
    public void onCallErrored(Throwable err) {}
  }

  static class MultiHeaderClientMiddlewareFactory implements FlightClientMiddleware.Factory {
    Map<String, List<byte[]>> lastBinaryHeaders = null;
    Map<String, List<String>> lastTextHeaders = null;

    @Override
    public FlightClientMiddleware onCallStarted(CallInfo info) {
      return new MultiHeaderClientMiddleware(this);
    }
  }

  static class MultiHeaderClientMiddleware implements FlightClientMiddleware {
    private final MultiHeaderClientMiddlewareFactory factory;

    public MultiHeaderClientMiddleware(MultiHeaderClientMiddlewareFactory factory) {
      this.factory = factory;
    }

    @Override
    public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) {
      for (final Map.Entry<String, List<byte[]>> entry : EXPECTED_BINARY_HEADERS.entrySet()) {
        entry.getValue().forEach((value) -> outgoingHeaders.insert(entry.getKey(), value));
        Assert.assertTrue(outgoingHeaders.containsKey(entry.getKey()));
      }
      for (final Map.Entry<String, List<String>> entry : EXPECTED_TEXT_HEADERS.entrySet()) {
        entry.getValue().forEach((value) -> outgoingHeaders.insert(entry.getKey(), value));
        Assert.assertTrue(outgoingHeaders.containsKey(entry.getKey()));
      }
    }

    @Override
    public void onHeadersReceived(CallHeaders incomingHeaders) {
      factory.lastBinaryHeaders = new HashMap<>();
      factory.lastTextHeaders = new HashMap<>();
      incomingHeaders.keys().forEach(header -> {
        if (header.endsWith("-bin")) {
          final List<byte[]> values = new ArrayList<>();
          incomingHeaders.getAllByte(header).forEach(values::add);
          factory.lastBinaryHeaders.put(header, values);
        } else {
          final List<String> values = new ArrayList<>();
          incomingHeaders.getAll(header).forEach(values::add);
          factory.lastTextHeaders.put(header, values);
        }
      });
    }

    @Override
    public void onCallCompleted(CallStatus status) {}
  }
}
