/*
 * Copyright 2023 The gRPC Authors
 *
 * Licensed 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 io.grpc.testing.integration;

import static java.util.concurrent.TimeUnit.SECONDS;
import static org.junit.Assert.assertTrue;

import io.grpc.ChannelCredentials;
import io.grpc.Grpc;
import io.grpc.InsecureChannelCredentials;
import io.grpc.ManagedChannel;
import io.grpc.alts.ComputeEngineChannelCredentials;
import java.util.ArrayList;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;

/**
 * Test client that can be used to verify that XDS federation works. A list of
 * server URIs (which can each be load balanced by different XDS servers), can
 * be configured via flags. A separate thread is created for each of these clients
 * and the configured test (either rpc_soak or channel_soak) is ran for each client
 * on each thread.
 */
public final class XdsFederationTestClient {
  private static final Logger logger =
      Logger.getLogger(XdsFederationTestClient.class.getName());

  /**
   * Entry point.
   */
  public static void main(String[] args) throws Exception {
    final XdsFederationTestClient client = new XdsFederationTestClient();
    client.parseArgs(args);
    client.setUp();
    client.run();
    System.exit(0);
  }

  private String serverUris = "";
  private String credentialsTypes = "";
  private int soakIterations = 10;
  private int soakMaxFailures = 0;
  private int soakPerIterationMaxAcceptableLatencyMs = 1000;
  private int soakOverallTimeoutSeconds = 10;
  private int soakMinTimeMsBetweenRpcs = 0;
  private int soakRequestSize = 271828;
  private int soakResponseSize = 314159;
  private String testCase = "rpc_soak";
  private final ArrayList<InnerClient> clients = new ArrayList<>();

  private void parseArgs(String[] args) {
    boolean usage = false;
    for (String arg : args) {
      if (!arg.startsWith("--")) {
        System.err.println("All arguments must start with '--': " + arg);
        usage = true;
        break;
      }
      String[] parts = arg.substring(2).split("=", 2);
      String key = parts[0];
      if (key.equals("help")) {
        usage = true;
        break;
      }
      if (parts.length != 2) {
        System.err.println("All arguments must be of the form --arg=value");
        usage = true;
        break;
      }
      String value = parts[1];
      switch (key) {
        case "server_uris":
          serverUris = value;
          break;
        case "credentials_types":
          credentialsTypes = value;
          break;
        case "test_case":
          testCase = value;
          break;
        case "soak_iterations":
          soakIterations = Integer.parseInt(value);
          break;
        case "soak_max_failures":
          soakMaxFailures = Integer.parseInt(value);
          break;
        case "soak_per_iteration_max_acceptable_latency_ms":
          soakPerIterationMaxAcceptableLatencyMs = Integer.parseInt(value);
          break;
        case "soak_overall_timeout_seconds":
          soakOverallTimeoutSeconds = Integer.parseInt(value);
          break;
        case "soak_min_time_ms_between_rpcs":
          soakMinTimeMsBetweenRpcs = Integer.parseInt(value);
          break;
        case "soak_request_size":
          soakRequestSize = Integer.parseInt(value);
          break;
        case "soak_response_size":
          soakResponseSize = Integer.parseInt(value);
          break;
        default:
          System.err.println("Unknown argument: " + key);
          usage = true;
          break;
      }
    }
    if (usage) {
      XdsFederationTestClient c = new XdsFederationTestClient();
      System.out.println(
          "Usage: [ARGS...]"
          + "\n"
          + "\n  --server_uris                         Comma separated list of server "
          + "URIs to make RPCs to. Default: "
          + c.serverUris
          + "\n  --credentials_types                   Comma-separated list of "
          + "\n                                        credentials, each entry is used "
          + "\n                                        for the server of the "
          + "\n                                        corresponding index in server_uris. "
          + "\n                                        Supported values: "
          + "compute_engine_channel_creds,INSECURE_CREDENTIALS. Default: "
          + c.credentialsTypes
          + "\n  --soak_iterations                     The number of iterations to use "
          + "\n                                        for the two tests: rpc_soak and "
          + "\n                                        channel_soak. Default: "
          + c.soakIterations
          + "\n  --soak_max_failures                   The number of iterations in soak "
          + "\n                                        tests that are allowed to fail "
          + "\n                                        (either due to non-OK status code "
          + "\n                                        or exceeding the per-iteration max "
          + "\n                                        acceptable latency). Default: "
          + c.soakMaxFailures
          + "\n  --soak_per_iteration_max_acceptable_latency_ms"
          + "\n                                        The number of milliseconds a "
          + "\n                                        single iteration in the two soak "
          + "\n                                        tests (rpc_soak and channel_soak) "
          + "\n                                        should take. Default: "
          + c.soakPerIterationMaxAcceptableLatencyMs
          + "\n  --soak_overall_timeout_seconds        The overall number of seconds "
          + "\n                                        after which a soak test should "
          + "\n                                        stop and fail, if the desired "
          + "\n                                        number of iterations have not yet "
          + "\n                                        completed. Default: "
          + c.soakOverallTimeoutSeconds
          + "\n  --soak_min_time_ms_between_rpcs       The minimum time in milliseconds "
          + "\n                                        between consecutive RPCs in a soak "
          + "\n                                        test (rpc_soak or channel_soak), "
          + "\n                                        useful for limiting QPS. Default: "
          + c.soakMinTimeMsBetweenRpcs
          + "\n  --test_case=TEST_CASE                 Test case to run. Valid options are:"
          + "\n      rpc_soak: sends --soak_iterations large_unary RPCs"
          + "\n      channel_soak: sends --soak_iterations RPCs, rebuilding the channel "
          + "each time."
          + "\n      Default: " + c.testCase
          + "\n --soak_request_size "
          + "\n                                        The request size in a soak RPC. Default "
          + c.soakRequestSize
          + "\n"
          + " --soak_response_size \n"
          + "                                          The response size in a soak RPC. Default"
          + " "
          + c.soakResponseSize
      );
      System.exit(1);
    }
  }

  void setUp() {
    String[] uris = serverUris.split(",", -1);
    String[] creds = credentialsTypes.split(",", -1);
    if (uris.length == 0) {
      throw new IllegalArgumentException("--server_uris is empty");
    }
    if (uris.length != creds.length) {
      throw new IllegalArgumentException("Number of entries in --server_uris "
          + "does not match number of entries in --credentials_types");
    }
    for (int i = 0; i < uris.length; i++) {
      clients.add(new InnerClient(creds[i], uris[i]));
    }
  }

  /**
   * Wraps a single client stub configuration and executes a
   * soak test case with that configuration.
   */
  class InnerClient {
    private final String credentialsType;
    private final String serverUri;
    private boolean runSucceeded = false;

    public InnerClient(String credentialsType, String serverUri) {
      this.credentialsType = credentialsType;
      this.serverUri = serverUri;
    }

    /**
     * Indicates whether run succeeded or not. This must only be called
     * after run() has finished.
     */
    public boolean runSucceeded() {
      return runSucceeded;
    }

    /**
     * Run the intended soak test.
     */
    public void run() throws InterruptedException {
      try {
        switch (testCase) {
          case "rpc_soak": {
            SoakClient.performSoakTest(
                serverUri,
                soakIterations,
                soakMaxFailures,
                soakPerIterationMaxAcceptableLatencyMs,
                soakMinTimeMsBetweenRpcs,
                soakOverallTimeoutSeconds,
                soakRequestSize,
                soakResponseSize,
                1,
                createChannel(),
                (currentChannel) -> currentChannel);
          }
              break;
          case "channel_soak": {
            SoakClient.performSoakTest(
                serverUri,
                soakIterations,
                soakMaxFailures,
                soakPerIterationMaxAcceptableLatencyMs,
                soakMinTimeMsBetweenRpcs,
                soakOverallTimeoutSeconds,
                soakRequestSize,
                soakResponseSize,
                1,
                createChannel(),
                (currentChannel) -> createNewChannel(currentChannel));
          }
            break;
          default:
            throw new RuntimeException("invalid testcase: " + testCase);
        }

        logger.info("Test case: " + testCase + " done for server: " + serverUri);
        runSucceeded = true;
      } catch (Exception e) {
        logger.info("Test case: " + testCase + " failed for server: " + serverUri);
        throw new RuntimeException(e);
      }
    }

    ManagedChannel createChannel() {
      ChannelCredentials channelCredentials;
      switch (credentialsType) {
        case "compute_engine_channel_creds":
          channelCredentials = ComputeEngineChannelCredentials.create();
          break;
        case "INSECURE_CREDENTIALS":
          channelCredentials = InsecureChannelCredentials.create();
          break;
        default:
          throw new IllegalArgumentException("Unknown custom credentials: " + credentialsType);
      }
      return Grpc.newChannelBuilder(serverUri, channelCredentials)
          .keepAliveTime(3600, SECONDS)
          .keepAliveTimeout(20, SECONDS)
          .build();
    }

    ManagedChannel createNewChannel(ManagedChannel currentChannel) {
      currentChannel.shutdownNow();
      try {
        currentChannel.awaitTermination(10, TimeUnit.SECONDS);
      } catch (InterruptedException e) {
        throw new RuntimeException("Interrupted while creating a new channel", e);
      }
      return createChannel();
    }
  }

  private void run() throws InterruptedException {
    logger.info("Begin test case: " + testCase);
    ArrayList<Thread> threads = new ArrayList<>();
    for (InnerClient c : clients) {
      Thread t = new Thread(() -> {
        try {
          c.run();
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt(); // Properly re-interrupt the thread
          throw new RuntimeException("Thread was interrupted during execution", e);
        }
      });
      t.start();
      threads.add(t);
    }
    for (Thread t : threads) {
      t.join();
    }
    for (InnerClient c : clients) {
      assertTrue(c.runSucceeded());
    }
    logger.info("Test case: " + testCase + " done for all clients!");
  }
}
