/*
 * 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.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

import com.google.common.base.Charsets;
import com.google.protobuf.ByteString;
import io.grpc.MethodDescriptor;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import org.apache.arrow.flight.FlightClient.ClientStreamListener;
import org.apache.arrow.flight.impl.Flight;
import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.VectorUnloader;
import org.apache.arrow.vector.ipc.WriteChannel;
import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
import org.apache.arrow.vector.ipc.message.IpcOption;
import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.DisabledOnOs;
import org.junit.jupiter.api.condition.OS;

/** Test the operations of a basic flight service. */
public class TestBasicOperation {

  @Test
  public void fastPathDefaults() {
    assertTrue(ArrowMessage.ENABLE_ZERO_COPY_READ);
    assertFalse(ArrowMessage.ENABLE_ZERO_COPY_WRITE);
  }

  @Test
  public void fallbackLocation() {
    assertEquals(
        "arrow-flight-reuse-connection://?", Location.reuseConnection().getUri().toString());
  }

  /** ARROW-6017: we should be able to construct locations for unknown schemes. */
  @Test
  public void unknownScheme() throws URISyntaxException {
    final Location location = new Location("s3://unknown");
    assertEquals("s3", location.getUri().getScheme());
  }

  @Test
  public void unknownSchemeRemote() throws Exception {
    test(
        c -> {
          try {
            final FlightInfo info = c.getInfo(FlightDescriptor.path("test"));
            assertEquals(
                new URI("https://example.com"),
                info.getEndpoints().get(0).getLocations().get(0).getUri());
          } catch (URISyntaxException e) {
            throw new RuntimeException(e);
          }
        });
  }

  @Test
  public void roundTripTicket() throws Exception {
    final Ticket ticket = new Ticket(new byte[] {0, 1, 2, 3, 4, 5});
    assertEquals(ticket, Ticket.deserialize(ticket.serialize()));
  }

  @Test
  public void roundTripInfo() throws Exception {
    final Map<String, String> metadata = new HashMap<>();
    metadata.put("foo", "bar");
    final Schema schema =
        new Schema(
            Arrays.asList(
                Field.nullable("a", new ArrowType.Int(32, true)),
                Field.nullable("b", new ArrowType.FixedSizeBinary(32))),
            metadata);
    final FlightInfo info1 =
        FlightInfo.builder(schema, FlightDescriptor.path(), Collections.emptyList())
            .setAppMetadata("foo".getBytes(StandardCharsets.UTF_8))
            .build();
    final FlightInfo info2 =
        new FlightInfo(
            schema,
            FlightDescriptor.command(new byte[2]),
            Collections.singletonList(
                FlightEndpoint.builder(
                        new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock"))
                    .setAppMetadata("bar".getBytes(StandardCharsets.UTF_8))
                    .build()),
            200,
            500);
    final FlightInfo info3 =
        new FlightInfo(
            schema,
            FlightDescriptor.path("a", "b"),
            Arrays.asList(
                new FlightEndpoint(
                    new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock")),
                new FlightEndpoint(
                    new Ticket(new byte[10]),
                    Location.forGrpcDomainSocket("/tmp/test.sock"),
                    forGrpcInsecure("localhost", 50051))),
            200,
            500);
    final FlightInfo info4 =
        new FlightInfo(
            schema,
            FlightDescriptor.path("a", "b"),
            Arrays.asList(
                new FlightEndpoint(
                    new Ticket(new byte[10]), Location.forGrpcDomainSocket("/tmp/test.sock")),
                new FlightEndpoint(
                    new Ticket(new byte[10]),
                    Location.forGrpcDomainSocket("/tmp/test.sock"),
                    forGrpcInsecure("localhost", 50051))),
            200,
            500, /*ordered*/
            true,
            IpcOption.DEFAULT);

    assertEquals(info1, FlightInfo.deserialize(info1.serialize()));
    assertEquals(info2, FlightInfo.deserialize(info2.serialize()));
    assertEquals(info3, FlightInfo.deserialize(info3.serialize()));
    assertEquals(info4, FlightInfo.deserialize(info4.serialize()));

    assertNotEquals(info3, info4);

    assertFalse(info1.getOrdered());
    assertFalse(info2.getOrdered());
    assertFalse(info3.getOrdered());
    assertTrue(info4.getOrdered());
  }

  @Test
  public void roundTripDescriptor() throws Exception {
    final FlightDescriptor cmd =
        FlightDescriptor.command("test command".getBytes(StandardCharsets.UTF_8));
    assertEquals(cmd, FlightDescriptor.deserialize(cmd.serialize()));
    final FlightDescriptor path = FlightDescriptor.path("foo", "bar", "test.arrow");
    assertEquals(path, FlightDescriptor.deserialize(path.serialize()));
  }

  @Test
  public void getDescriptors() throws Exception {
    test(
        c -> {
          int count = 0;
          for (FlightInfo unused : c.listFlights(Criteria.ALL)) {
            count += 1;
          }
          assertEquals(1, count);
        });
  }

  @Test
  public void getDescriptorsWithCriteria() throws Exception {
    test(
        c -> {
          int count = 0;
          for (FlightInfo unused : c.listFlights(new Criteria(new byte[] {1}))) {

            count += 1;
          }
          assertEquals(0, count);
        });
  }

  @Test
  public void getDescriptor() throws Exception {
    test(
        c -> {
          System.out.println(c.getInfo(FlightDescriptor.path("hello")).getDescriptor());
        });
  }

  @Test
  public void getSchema() throws Exception {
    test(
        c -> {
          System.out.println(c.getSchema(FlightDescriptor.path("hello")).getSchema());
        });
  }

  @Test
  public void listActions() throws Exception {
    test(
        c -> {
          for (ActionType at : c.listActions()) {
            System.out.println(at.getType());
          }
        });
  }

  @Test
  public void doAction() throws Exception {
    test(
        c -> {
          Iterator<Result> stream = c.doAction(new Action("hello"));

          assertTrue(stream.hasNext());
          Result r = stream.next();
          assertArrayEquals("world".getBytes(Charsets.UTF_8), r.getBody());
        });
    test(
        c -> {
          Iterator<Result> stream = c.doAction(new Action("hellooo"));

          assertTrue(stream.hasNext());
          Result r = stream.next();
          assertArrayEquals("world".getBytes(Charsets.UTF_8), r.getBody());

          assertTrue(stream.hasNext());
          r = stream.next();
          assertArrayEquals("!".getBytes(Charsets.UTF_8), r.getBody());
          assertFalse(stream.hasNext());
        });
  }

  @Test
  public void putStream() throws Exception {
    test(
        (c, a) -> {
          final int size = 10;

          IntVector iv = new IntVector("c1", a);

          try (VectorSchemaRoot root = VectorSchemaRoot.of(iv)) {
            ClientStreamListener listener =
                c.startPut(FlightDescriptor.path("hello"), root, new AsyncPutListener());

            // batch 1
            root.allocateNew();
            for (int i = 0; i < size; i++) {
              iv.set(i, i);
            }
            iv.setValueCount(size);
            root.setRowCount(size);
            listener.putNext();

            // batch 2

            root.allocateNew();
            for (int i = 0; i < size; i++) {
              iv.set(i, i + size);
            }
            iv.setValueCount(size);
            root.setRowCount(size);
            listener.putNext();
            root.clear();
            listener.completed();

            // wait for ack to avoid memory leaks.
            listener.getResult();
          }
        });
  }

  @Test
  public void propagateErrors() throws Exception {
    test(
        client -> {
          FlightTestUtil.assertCode(
              FlightStatusCode.UNIMPLEMENTED,
              () -> {
                client.doAction(new Action("invalid-action")).forEachRemaining(action -> fail());
              });
        });
  }

  @Test
  public void getStream() throws Exception {
    test(
        c -> {
          try (final 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++;
              }
            }
          } catch (Exception e) {
            throw new RuntimeException(e);
          }
        });
  }

  /** Ensure the client is configured to accept large messages. */
  @Test
  @DisabledOnOs(
      value = {OS.WINDOWS},
      disabledReason = "https://github.com/apache/arrow/issues/33237: flaky test")
  public void getStreamLargeBatch() throws Exception {
    test(
        c -> {
          try (final FlightStream stream = c.getStream(new Ticket(Producer.TICKET_LARGE_BATCH))) {
            assertEquals(128, stream.getRoot().getFieldVectors().size());
            assertTrue(stream.next());
            assertEquals(65536, stream.getRoot().getRowCount());
            assertTrue(stream.next());
            assertEquals(65536, stream.getRoot().getRowCount());
            assertFalse(stream.next());
          } catch (Exception e) {
            throw new RuntimeException(e);
          }
        });
  }

  /** Ensure the server is configured to accept large messages. */
  @Test
  public void startPutLargeBatch() throws Exception {
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
      final List<FieldVector> vectors = new ArrayList<>();
      for (int col = 0; col < 128; col++) {
        final BigIntVector vector = new BigIntVector("f" + col, allocator);
        for (int row = 0; row < 65536; row++) {
          vector.setSafe(row, row);
        }
        vectors.add(vector);
      }
      test(
          c -> {
            try (final VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
              root.setRowCount(65536);
              final ClientStreamListener stream =
                  c.startPut(FlightDescriptor.path(""), root, new SyncPutListener());
              stream.putNext();
              stream.putNext();
              stream.completed();
              stream.getResult();
            } catch (Exception e) {
              throw new RuntimeException(e);
            }
          });
    }
  }

  private void test(Consumer<FlightClient> consumer) throws Exception {
    test(
        (c, a) -> {
          consumer.accept(c);
        });
  }

  private void test(BiConsumer<FlightClient, BufferAllocator> consumer) throws Exception {
    try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
        Producer producer = new Producer(a);
        FlightServer s =
            FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0), producer).build().start()) {

      try (FlightClient c = FlightClient.builder(a, s.getLocation()).build()) {
        try (BufferAllocator testAllocator = a.newChildAllocator("testcase", 0, Long.MAX_VALUE)) {
          consumer.accept(c, testAllocator);
        }
      }
    }
  }

  /** Helper method to convert an ArrowMessage into a Protobuf message. */
  private Flight.FlightData arrowMessageToProtobuf(
      MethodDescriptor.Marshaller<ArrowMessage> marshaller, ArrowMessage message)
      throws IOException {
    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (final InputStream serialized = marshaller.stream(message)) {
      final byte[] buf = new byte[1024];
      while (true) {
        int read = serialized.read(buf);
        if (read < 0) {
          break;
        }
        baos.write(buf, 0, read);
      }
    }
    final byte[] serializedMessage = baos.toByteArray();
    return Flight.FlightData.parseFrom(serializedMessage);
  }

  /**
   * ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields).
   */
  @Test
  public void testProtobufRecordBatchCompatibility() throws Exception {
    final Schema schema =
        new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true))));
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
        final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
      final VectorUnloader unloader = new VectorUnloader(root);
      root.setRowCount(0);
      final MethodDescriptor.Marshaller<ArrowMessage> marshaller =
          ArrowMessage.createMarshaller(allocator);
      try (final ArrowMessage message =
          new ArrowMessage(
              unloader.getRecordBatch(), /* appMetadata */
              null, /* tryZeroCopy */
              false,
              IpcOption.DEFAULT)) {
        assertEquals(ArrowMessage.HeaderType.RECORD_BATCH, message.getMessageType());
        // Should have at least one empty body buffer (there may be multiple for e.g. data and
        // validity)
        Iterator<ArrowBuf> iterator = message.getBufs().iterator();
        assertTrue(iterator.hasNext());
        while (iterator.hasNext()) {
          assertEquals(0, iterator.next().capacity());
        }
        final Flight.FlightData protobufData =
            arrowMessageToProtobuf(marshaller, message).toBuilder().clearDataBody().build();
        assertEquals(0, protobufData.getDataBody().size());
        ArrowMessage parsedMessage =
            marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray()));
        // Should have an empty body buffer
        Iterator<ArrowBuf> parsedIterator = parsedMessage.getBufs().iterator();
        assertTrue(parsedIterator.hasNext());
        assertEquals(0, parsedIterator.next().capacity());
        // Should have only one (the parser synthesizes exactly one); in the case of empty buffers,
        // this is equivalent
        assertFalse(parsedIterator.hasNext());
        // Should not throw
        final ArrowRecordBatch rb = parsedMessage.asRecordBatch();
        assertEquals(rb.computeBodyLength(), 0);
      }
    }
  }

  /**
   * ARROW-10962: accept FlightData messages generated by Protobuf (which can omit empty fields).
   */
  @Test
  public void testProtobufSchemaCompatibility() throws Exception {
    final Schema schema =
        new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true))));
    try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
      final MethodDescriptor.Marshaller<ArrowMessage> marshaller =
          ArrowMessage.createMarshaller(allocator);
      Flight.FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]).toProtocol();
      try (final ArrowMessage message = new ArrowMessage(descriptor, schema, IpcOption.DEFAULT)) {
        assertEquals(ArrowMessage.HeaderType.SCHEMA, message.getMessageType());
        // Should have no body buffers
        assertFalse(message.getBufs().iterator().hasNext());
        final Flight.FlightData protobufData =
            arrowMessageToProtobuf(marshaller, message).toBuilder()
                .setDataBody(ByteString.EMPTY)
                .build();
        assertEquals(0, protobufData.getDataBody().size());
        final ArrowMessage parsedMessage =
            marshaller.parse(new ByteArrayInputStream(protobufData.toByteArray()));
        // Should have no body buffers
        assertFalse(parsedMessage.getBufs().iterator().hasNext());
        // Should not throw
        parsedMessage.asSchema();
      }
    }
  }

  @Test
  public void testGrpcInsecureLocation() throws Exception {
    Location location = Location.forGrpcInsecure(LOCALHOST, 9000);
    assertEquals(
        new URI(LocationSchemes.GRPC_INSECURE, null, LOCALHOST, 9000, null, null, null),
        location.getUri());
    assertEquals(new InetSocketAddress(LOCALHOST, 9000), location.toSocketAddress());
  }

  @Test
  public void testGrpcTlsLocation() throws Exception {
    Location location = Location.forGrpcTls(LOCALHOST, 9000);
    assertEquals(
        new URI(LocationSchemes.GRPC_TLS, null, LOCALHOST, 9000, null, null, null),
        location.getUri());
    assertEquals(new InetSocketAddress(LOCALHOST, 9000), location.toSocketAddress());
  }

  /** An example FlightProducer for test purposes. */
  public static class Producer implements FlightProducer, AutoCloseable {
    static final byte[] TICKET_LARGE_BATCH = "large-batch".getBytes(StandardCharsets.UTF_8);

    private final BufferAllocator allocator;

    public Producer(BufferAllocator allocator) {
      super();
      this.allocator = allocator;
    }

    @Override
    public void listFlights(
        CallContext context, Criteria criteria, StreamListener<FlightInfo> listener) {
      if (criteria.getExpression().length > 0) {
        // Don't send anything if criteria are set
        listener.onCompleted();
      }

      Flight.FlightInfo getInfo =
          Flight.FlightInfo.newBuilder()
              .setFlightDescriptor(
                  Flight.FlightDescriptor.newBuilder()
                      .setType(DescriptorType.CMD)
                      .setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8)))
              .build();
      try {
        listener.onNext(new FlightInfo(getInfo));
      } catch (URISyntaxException e) {
        listener.onError(e);
        return;
      }
      listener.onCompleted();
    }

    @Override
    public Runnable acceptPut(
        CallContext context, FlightStream flightStream, StreamListener<PutResult> ackStream) {
      return () -> {
        while (flightStream.next()) {
          // Drain the stream
        }
      };
    }

    @Override
    public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
      if (Arrays.equals(TICKET_LARGE_BATCH, ticket.getBytes())) {
        getLargeBatch(listener);
        return;
      }
      final int size = 10;

      IntVector iv = new IntVector("c1", allocator);
      VectorSchemaRoot root = VectorSchemaRoot.of(iv);
      listener.start(root);

      // batch 1
      root.allocateNew();
      for (int i = 0; i < size; i++) {
        iv.set(i, i);
      }
      iv.setValueCount(size);
      root.setRowCount(size);
      listener.putNext();

      // batch 2

      root.allocateNew();
      for (int i = 0; i < size; i++) {
        iv.set(i, i + size);
      }
      iv.setValueCount(size);
      root.setRowCount(size);
      listener.putNext();
      root.clear();
      listener.completed();
    }

    private void getLargeBatch(ServerStreamListener listener) {
      final List<FieldVector> vectors = new ArrayList<>();
      for (int col = 0; col < 128; col++) {
        final BigIntVector vector = new BigIntVector("f" + col, allocator);
        for (int row = 0; row < 65536; row++) {
          vector.setSafe(row, row);
        }
        vectors.add(vector);
      }
      try (final VectorSchemaRoot root = new VectorSchemaRoot(vectors)) {
        root.setRowCount(65536);
        listener.start(root);
        listener.putNext();
        listener.putNext();
        listener.completed();
      }
    }

    @Override
    public void close() throws Exception {
      allocator.close();
    }

    @Override
    public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
      try {
        Flight.FlightInfo getInfo =
            Flight.FlightInfo.newBuilder()
                .setSchema(schemaToByteString(new Schema(Collections.emptyList())))
                .setFlightDescriptor(
                    Flight.FlightDescriptor.newBuilder()
                        .setType(DescriptorType.CMD)
                        .setCmd(ByteString.copyFrom("cool thing", Charsets.UTF_8)))
                .addEndpoint(
                    Flight.FlightEndpoint.newBuilder()
                        .addLocation(new Location("https://example.com").toProtocol()))
                .build();
        return new FlightInfo(getInfo);
      } catch (URISyntaxException e) {
        throw new RuntimeException(e);
      }
    }

    private static ByteString schemaToByteString(Schema schema) {
      try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
        MessageSerializer.serialize(
            new WriteChannel(Channels.newChannel(baos)), schema, IpcOption.DEFAULT);
        return ByteString.copyFrom(baos.toByteArray());
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    }

    @Override
    public void doAction(CallContext context, Action action, StreamListener<Result> listener) {
      switch (action.getType()) {
        case "hello":
          {
            listener.onNext(new Result("world".getBytes(Charsets.UTF_8)));
            listener.onCompleted();
            break;
          }
        case "hellooo":
          {
            listener.onNext(new Result("world".getBytes(Charsets.UTF_8)));
            listener.onNext(new Result("!".getBytes(Charsets.UTF_8)));
            listener.onCompleted();
            break;
          }
        default:
          listener.onError(
              CallStatus.UNIMPLEMENTED
                  .withDescription("Action not implemented: " + action.getType())
                  .toRuntimeException());
      }
    }

    @Override
    public void listActions(CallContext context, StreamListener<ActionType> listener) {
      listener.onNext(new ActionType("get", ""));
      listener.onNext(new ActionType("put", ""));
      listener.onNext(new ActionType("hello", ""));
      listener.onCompleted();
    }
  }
}
