/*
 * 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.beam.sdk.extensions.gcp.options;

import com.google.api.client.googleapis.services.AbstractGoogleClient;
import com.google.api.client.googleapis.services.AbstractGoogleClientRequest;
import com.google.api.client.googleapis.services.GoogleClientRequestInitializer;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import org.apache.beam.sdk.options.Description;
import org.apache.beam.sdk.options.PipelineOptions;

/**
 * These options configure debug settings for Google API clients created within the Apache Beam SDK.
 */
public interface GoogleApiDebugOptions extends PipelineOptions {
  /**
   * This option enables tracing of API calls to Google services used within the Apache Beam SDK.
   * Values are expected in JSON format <code>{"ApiName":"TraceDestination",...}
   * </code> where the {@code ApiName} represents the request classes canonical name. The {@code
   * TraceDestination} is a logical trace consumer to whom the trace will be reported. Typically,
   * "producer" is the right destination to use: this makes API traces available to the team
   * offering the API. Note that by enabling this option, the contents of the requests to and from
   * Google Cloud services will be made available to Google. For example, by specifying <code>
   * {"Dataflow":"producer"}</code>, all calls to the Dataflow service will be made available to
   * Google, specifically to the Google Cloud Dataflow team.
   */
  @Description(
      "This option enables tracing of API calls to Google services used within the Apache "
          + "Beam SDK. Values are expected in JSON format {\"ApiName\":\"TraceDestination\",...} "
          + "where the ApiName represents the request classes canonical name. The TraceDestination is "
          + "a logical trace consumer to whom the trace will be reported. Typically, \"producer\" is "
          + "the right destination to use: this makes API traces available to the team offering the "
          + "API. Note that by enabling this option, the contents of the requests to and from "
          + "Google Cloud services will be made available to Google. For example, by specifying "
          + "{\"Dataflow\":\"producer\"}, all calls to the Dataflow service will be made available to "
          + "Google, specifically to the Google Cloud Dataflow team.")
  GoogleApiTracer getGoogleApiTrace();

  void setGoogleApiTrace(GoogleApiTracer commands);

  /**
   * A {@link GoogleClientRequestInitializer} that adds the trace destination to Google API calls.
   */
  class GoogleApiTracer extends HashMap<String, String> implements GoogleClientRequestInitializer {
    /**
     * Creates a {@link GoogleApiTracer} that sets the trace destination on all calls that match the
     * given client type.
     */
    public GoogleApiTracer addTraceFor(AbstractGoogleClient client, String traceDestination) {
      put(client.getClass().getCanonicalName(), traceDestination);
      return this;
    }

    /**
     * Creates a {@link GoogleApiTracer} that sets the trace {@code traceDestination} on all calls
     * that match for the given request type.
     */
    public GoogleApiTracer addTraceFor(
        AbstractGoogleClientRequest<?> request, String traceDestination) {
      put(request.getClass().getCanonicalName(), traceDestination);
      return this;
    }

    @Override
    public void initialize(AbstractGoogleClientRequest<?> request) throws IOException {
      for (Map.Entry<String, String> entry : this.entrySet()) {
        if (request.getClass().getCanonicalName().contains(entry.getKey())) {
          request.set("$trace", entry.getValue());
        }
      }
    }
  }
}