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

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

import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import java.io.IOException;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.Criteria;
import org.apache.arrow.flight.FlightClient;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightProducer;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.FlightTestUtil;
import org.apache.arrow.flight.NoOpFlightProducer;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.flight.grpc.CredentialCallOption;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

public class TestBasicAuth2 {

  private static final String USERNAME_1 = "flight1";
  private static final String USERNAME_2 = "flight2";
  private static final String NO_USERNAME = "";
  private static final String PASSWORD_1 = "woohoo1";
  private static final String PASSWORD_2 = "woohoo2";
  private static BufferAllocator allocator;
  private static FlightServer server;
  private static FlightClient client;
  private static FlightClient client2;

  @BeforeAll
  public static void setup() throws Exception {
    allocator = new RootAllocator(Long.MAX_VALUE);
    startServerAndClient();
  }

  private static FlightProducer getFlightProducer() {
    return new NoOpFlightProducer() {
      @Override
      public void listFlights(
          CallContext context, Criteria criteria, StreamListener<FlightInfo> listener) {
        if (!context.peerIdentity().equals(USERNAME_1)
            && !context.peerIdentity().equals(USERNAME_2)) {
          listener.onError(new IllegalArgumentException("Invalid username"));
          return;
        }
        listener.onCompleted();
      }

      @Override
      public void getStream(CallContext context, Ticket ticket, ServerStreamListener listener) {
        if (!context.peerIdentity().equals(USERNAME_1)
            && !context.peerIdentity().equals(USERNAME_2)) {
          listener.error(new IllegalArgumentException("Invalid username"));
          return;
        }
        final Schema pojoSchema =
            new Schema(ImmutableList.of(Field.nullable("a", Types.MinorType.BIGINT.getType())));
        try (VectorSchemaRoot root = VectorSchemaRoot.create(pojoSchema, allocator)) {
          listener.start(root);
          root.allocateNew();
          root.setRowCount(4095);
          listener.putNext();
          listener.completed();
        }
      }
    };
  }

  private static void startServerAndClient() throws IOException {
    final FlightProducer flightProducer = getFlightProducer();
    server =
        FlightServer.builder(allocator, forGrpcInsecure(LOCALHOST, 0), flightProducer)
            .headerAuthenticator(
                new GeneratedBearerTokenAuthenticator(
                    new BasicCallHeaderAuthenticator(TestBasicAuth2::validate)))
            .build()
            .start();
    client = FlightClient.builder(allocator, server.getLocation()).build();
  }

  @AfterAll
  public static void shutdown() throws Exception {
    AutoCloseables.close(client, client2, server);
    client = null;
    client2 = null;
    server = null;

    allocator.getChildAllocators().forEach(BufferAllocator::close);
    AutoCloseables.close(allocator);
    allocator = null;
  }

  private void startClient2() throws IOException {
    client2 = FlightClient.builder(allocator, server.getLocation()).build();
  }

  private static CallHeaderAuthenticator.AuthResult validate(String username, String password) {
    if (Strings.isNullOrEmpty(username)) {
      throw CallStatus.UNAUTHENTICATED
          .withDescription("Credentials not supplied.")
          .toRuntimeException();
    }
    final String identity;
    if (USERNAME_1.equals(username) && PASSWORD_1.equals(password)) {
      identity = USERNAME_1;
    } else if (USERNAME_2.equals(username) && PASSWORD_2.equals(password)) {
      identity = USERNAME_2;
    } else {
      throw CallStatus.UNAUTHENTICATED
          .withDescription("Username or password is invalid.")
          .toRuntimeException();
    }
    return () -> identity;
  }

  @Test
  public void validAuthWithBearerAuthServer() throws IOException {
    testValidAuth(client);
  }

  @Test
  public void validAuthWithMultipleClientsWithSameCredentialsWithBearerAuthServer()
      throws IOException {
    startClient2();
    testValidAuthWithMultipleClientsWithSameCredentials(client, client2);
  }

  @Test
  public void validAuthWithMultipleClientsWithDifferentCredentialsWithBearerAuthServer()
      throws IOException {
    startClient2();
    testValidAuthWithMultipleClientsWithDifferentCredentials(client, client2);
  }

  @Test
  public void asyncCall() throws Exception {
    final CredentialCallOption bearerToken =
        client.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
    client.listFlights(Criteria.ALL, bearerToken);
    try (final FlightStream s = client.getStream(new Ticket(new byte[1]), bearerToken)) {
      while (s.next()) {
        assertEquals(4095, s.getRoot().getRowCount());
      }
    }
  }

  @Test
  public void invalidAuthWithBearerAuthServer() throws IOException {
    testInvalidAuth(client);
  }

  @Test
  public void didntAuthWithBearerAuthServer() throws IOException {
    didntAuth(client);
  }

  private void testValidAuth(FlightClient client) {
    final CredentialCallOption bearerToken =
        client.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
    assertTrue(ImmutableList.copyOf(client.listFlights(Criteria.ALL, bearerToken)).isEmpty());
  }

  private void testValidAuthWithMultipleClientsWithSameCredentials(
      FlightClient client1, FlightClient client2) {
    final CredentialCallOption bearerToken1 =
        client1.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
    final CredentialCallOption bearerToken2 =
        client2.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
    assertTrue(ImmutableList.copyOf(client1.listFlights(Criteria.ALL, bearerToken1)).isEmpty());
    assertTrue(ImmutableList.copyOf(client2.listFlights(Criteria.ALL, bearerToken2)).isEmpty());
  }

  private void testValidAuthWithMultipleClientsWithDifferentCredentials(
      FlightClient client1, FlightClient client2) {
    final CredentialCallOption bearerToken1 =
        client1.authenticateBasicToken(USERNAME_1, PASSWORD_1).get();
    final CredentialCallOption bearerToken2 =
        client2.authenticateBasicToken(USERNAME_2, PASSWORD_2).get();
    assertTrue(ImmutableList.copyOf(client1.listFlights(Criteria.ALL, bearerToken1)).isEmpty());
    assertTrue(ImmutableList.copyOf(client2.listFlights(Criteria.ALL, bearerToken2)).isEmpty());
  }

  private void testInvalidAuth(FlightClient client) {
    FlightTestUtil.assertCode(
        FlightStatusCode.UNAUTHENTICATED, () -> client.authenticateBasicToken(USERNAME_1, "WRONG"));

    FlightTestUtil.assertCode(
        FlightStatusCode.UNAUTHENTICATED,
        () -> client.authenticateBasicToken(NO_USERNAME, PASSWORD_1));

    FlightTestUtil.assertCode(
        FlightStatusCode.UNAUTHENTICATED,
        () -> client.listFlights(Criteria.ALL).forEach(action -> fail()));
  }

  private void didntAuth(FlightClient client) {
    FlightTestUtil.assertCode(
        FlightStatusCode.UNAUTHENTICATED,
        () -> client.listFlights(Criteria.ALL).forEach(action -> fail()));
  }
}
