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

import com.google.protobuf.Any;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.rpc.Status;
import io.grpc.Metadata;
import io.grpc.StatusRuntimeException;
import io.grpc.protobuf.ProtoUtils;
import io.grpc.protobuf.StatusProto;
import java.nio.charset.StandardCharsets;
import org.apache.arrow.flight.perf.impl.PerfOuterClass;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.junit.jupiter.api.Test;

public class TestErrorMetadata {
  private static final Metadata.BinaryMarshaller<Status> marshaller =
      ProtoUtils.metadataMarshaller(Status.getDefaultInstance());

  /** Ensure metadata attached to a gRPC error is propagated. */
  @Test
  public void testGrpcMetadata() throws Exception {
    PerfOuterClass.Perf perf =
        PerfOuterClass.Perf.newBuilder()
            .setStreamCount(12)
            .setRecordsPerBatch(1000)
            .setRecordsPerStream(1000000L)
            .build();
    StatusRuntimeExceptionProducer producer = new StatusRuntimeExceptionProducer(perf);
    try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        final FlightServer s =
            FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), producer)
                .build()
                .start();
        final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) {
      final CallStatus flightStatus =
          FlightTestUtil.assertCode(
              FlightStatusCode.CANCELLED,
              () -> {
                FlightStream stream =
                    client.getStream(new Ticket("abs".getBytes(StandardCharsets.UTF_8)));
                stream.next();
              });
      PerfOuterClass.Perf newPerf = null;
      ErrorFlightMetadata metadata = flightStatus.metadata();
      assertNotNull(metadata);
      assertEquals(2, metadata.keys().size());
      assertTrue(metadata.containsKey("grpc-status-details-bin"));
      Status status = marshaller.parseBytes(metadata.getByte("grpc-status-details-bin"));
      for (Any details : status.getDetailsList()) {
        if (details.is(PerfOuterClass.Perf.class)) {
          try {
            newPerf = details.unpack(PerfOuterClass.Perf.class);
          } catch (InvalidProtocolBufferException e) {
            fail();
          }
        }
      }
      assertNotNull(newPerf);
      assertEquals(perf, newPerf);
    }
  }

  /** Ensure metadata attached to a Flight error is propagated. */
  @Test
  public void testFlightMetadata() throws Exception {
    try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        final FlightServer s =
            FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), new CallStatusProducer())
                .build()
                .start();
        final FlightClient client = FlightClient.builder(allocator, s.getLocation()).build()) {
      CallStatus flightStatus =
          FlightTestUtil.assertCode(
              FlightStatusCode.INVALID_ARGUMENT,
              () -> {
                FlightStream stream = client.getStream(new Ticket(new byte[0]));
                stream.next();
              });
      ErrorFlightMetadata metadata = flightStatus.metadata();
      assertNotNull(metadata);
      assertEquals("foo", metadata.get("x-foo"));
      assertArrayEquals(new byte[] {1}, metadata.getByte("x-bar-bin"));

      flightStatus =
          FlightTestUtil.assertCode(
              FlightStatusCode.INVALID_ARGUMENT,
              () -> {
                client.getInfo(FlightDescriptor.command(new byte[0]));
              });
      metadata = flightStatus.metadata();
      assertNotNull(metadata);
      assertEquals("foo", metadata.get("x-foo"));
      assertArrayEquals(new byte[] {1}, metadata.getByte("x-bar-bin"));
    }
  }

  private static class StatusRuntimeExceptionProducer extends NoOpFlightProducer {
    private final PerfOuterClass.Perf perf;

    private StatusRuntimeExceptionProducer(PerfOuterClass.Perf perf) {
      this.perf = perf;
    }

    @Override
    public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
      StatusRuntimeException sre =
          StatusProto.toStatusRuntimeException(
              Status.newBuilder()
                  .setCode(1)
                  .setMessage("Testing 1 2 3")
                  .addDetails(Any.pack(perf, "arrow/meta/types"))
                  .build());
      listener.error(sre);
    }
  }

  private static class CallStatusProducer extends NoOpFlightProducer {
    ErrorFlightMetadata metadata;

    CallStatusProducer() {
      this.metadata = new ErrorFlightMetadata();
      metadata.insert("x-foo", "foo");
      metadata.insert("x-bar-bin", new byte[] {1});
    }

    @Override
    public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
      listener.error(
          CallStatus.INVALID_ARGUMENT
              .withDescription("Failed")
              .withMetadata(metadata)
              .toRuntimeException());
    }

    @Override
    public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
      throw CallStatus.INVALID_ARGUMENT
          .withDescription("Failed")
          .withMetadata(metadata)
          .toRuntimeException();
    }
  }
}
