/*
 * 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.runners.spark.metrics;

import com.google.common.base.Function;
import com.google.common.base.Objects;
import com.google.common.base.Predicate;
import com.google.common.collect.FluentIterable;
import java.util.Set;
import org.apache.beam.sdk.metrics.DistributionData;
import org.apache.beam.sdk.metrics.DistributionResult;
import org.apache.beam.sdk.metrics.GaugeData;
import org.apache.beam.sdk.metrics.GaugeResult;
import org.apache.beam.sdk.metrics.MetricKey;
import org.apache.beam.sdk.metrics.MetricName;
import org.apache.beam.sdk.metrics.MetricNameFilter;
import org.apache.beam.sdk.metrics.MetricQueryResults;
import org.apache.beam.sdk.metrics.MetricResult;
import org.apache.beam.sdk.metrics.MetricResults;
import org.apache.beam.sdk.metrics.MetricUpdates.MetricUpdate;
import org.apache.beam.sdk.metrics.MetricsFilter;


/**
 * Implementation of {@link MetricResults} for the Spark Runner.
 */
public class SparkMetricResults extends MetricResults {

  @Override
  public MetricQueryResults queryMetrics(MetricsFilter filter) {
    return new SparkMetricQueryResults(filter);
  }

  private static class SparkMetricQueryResults implements MetricQueryResults {
    private final MetricsFilter filter;

    SparkMetricQueryResults(MetricsFilter filter) {
      this.filter = filter;
    }

    @Override
    public Iterable<MetricResult<Long>> counters() {
      return
          FluentIterable
              .from(SparkMetricsContainer.getCounters())
              .filter(matchesFilter(filter))
              .transform(TO_COUNTER_RESULT)
              .toList();
    }

    @Override
    public Iterable<MetricResult<DistributionResult>> distributions() {
      return
          FluentIterable
              .from(SparkMetricsContainer.getDistributions())
              .filter(matchesFilter(filter))
              .transform(TO_DISTRIBUTION_RESULT)
              .toList();
    }

    @Override
    public Iterable<MetricResult<GaugeResult>> gauges() {
      return
          FluentIterable
              .from(SparkMetricsContainer.getGauges())
              .filter(matchesFilter(filter))
              .transform(TO_GAUGE_RESULT)
              .toList();
    }

    private Predicate<MetricUpdate<?>> matchesFilter(final MetricsFilter filter) {
      return new Predicate<MetricUpdate<?>>() {
        @Override
        public boolean apply(MetricUpdate<?> metricResult) {
          return matches(filter, metricResult.getKey());
        }
      };
    }

    private boolean matches(MetricsFilter filter, MetricKey key) {
      return matchesName(key.metricName(), filter.names())
          && matchesScope(key.stepName(), filter.steps());
    }

    private boolean matchesName(MetricName metricName, Set<MetricNameFilter> nameFilters) {
      if (nameFilters.isEmpty()) {
        return true;
      }

      for (MetricNameFilter nameFilter : nameFilters) {
        if ((nameFilter.getName() == null || nameFilter.getName().equals(metricName.name()))
            && Objects.equal(metricName.namespace(), nameFilter.getNamespace())) {
          return true;
        }
      }

      return false;
    }

    private boolean matchesScope(String actualScope, Set<String> scopes) {
      if (scopes.isEmpty() || scopes.contains(actualScope)) {
        return true;
      }

      for (String scope : scopes) {
        if (actualScope.startsWith(scope)) {
          return true;
        }
      }

      return false;
    }
  }

  private static final Function<MetricUpdate<DistributionData>, MetricResult<DistributionResult>>
      TO_DISTRIBUTION_RESULT =
      new Function<MetricUpdate<DistributionData>, MetricResult<DistributionResult>>() {
        @Override
        public MetricResult<DistributionResult> apply(MetricUpdate<DistributionData> metricResult) {
          if (metricResult != null) {
            MetricKey key = metricResult.getKey();
            return new SparkMetricResult<>(key.metricName(), key.stepName(),
                metricResult.getUpdate().extractResult());
          } else {
            return null;
          }
        }
      };

  private static final Function<MetricUpdate<Long>, MetricResult<Long>>
      TO_COUNTER_RESULT =
      new Function<MetricUpdate<Long>, MetricResult<Long>>() {
        @Override
        public MetricResult<Long> apply(MetricUpdate<Long> metricResult) {
          if (metricResult != null) {
            MetricKey key = metricResult.getKey();
            return new SparkMetricResult<>(key.metricName(), key.stepName(),
                metricResult.getUpdate());
          } else {
            return null;
          }
        }
      };

  private static final Function<MetricUpdate<GaugeData>, MetricResult<GaugeResult>>
      TO_GAUGE_RESULT =
      new Function<MetricUpdate<GaugeData>, MetricResult<GaugeResult>>() {
        @Override
        public MetricResult<GaugeResult> apply(MetricUpdate<GaugeData> metricResult) {
          if (metricResult != null) {
            MetricKey key = metricResult.getKey();
            return new SparkMetricResult<>(key.metricName(), key.stepName(),
                metricResult.getUpdate().extractResult());
          } else {
            return null;
          }
        }
      };

  private static class SparkMetricResult<T> implements MetricResult<T> {
    private final MetricName name;
    private final String step;
    private final T result;

    SparkMetricResult(MetricName name, String step, T result) {
      this.name = name;
      this.step = step;
      this.result = result;
    }

    @Override
    public MetricName name() {
      return name;
    }

    @Override
    public String step() {
      return step;
    }

    @Override
    public T committed() {
      throw new UnsupportedOperationException("Spark runner does not currently support committed"
          + " metrics results. Please use 'attempted' instead.");
    }

    @Override
    public T attempted() {
      return result;
    }
  }
}

@@ -892,47 +892,8 @@
 on;%0A
-import com.google.common.base.Objects;%0A
 impo
@@ -982,30 +982,8 @@
 le;%0A
-import java.util.Set;%0A
 impo
@@ -1221,19 +1221,25 @@
 s.Metric
-Key
+Filtering
 ;%0Aimport
@@ -1273,20 +1273,19 @@
 s.Metric
-Name
+Key
 ;%0Aimport
@@ -1323,22 +1323,16 @@
 tricName
-Filter
 ;%0Aimport
@@ -3171,1032 +3171,81 @@
 urn 
-matches(filter, metricResult.getKey());%0A        %7D%0A      %7D;%0A    %7D%0A%0A    private boolean matches(
 Metric
-s
 Filter
- filter, MetricKey key) %7B%0A      return matchesName(key.metricName(), filter.names())%0A          && matchesScope(key.stepName(), filter.steps());%0A    %7D%0A%0A    private boolean matchesName(MetricName metricName, Set%3CMetricNameFilter%3E nameFilters) %7B%0A      if (nameFilters.isEmpty()) %7B%0A        return true;%0A      %7D%0A%0A      for (MetricNameFilter nameFilter : nameFilters) %7B%0A        if ((nameFilter.getName() == null %7C%7C nameFilter.getName().equals(metricName.name()))%0A            && Objects.equal(metricName.namespace(), nameFilter.getNamespace())) %7B%0A          return true;%0A        %7D%0A      %7D%0A%0A      return false;%0A    %7D%0A%0A    private boolean matchesScope(String actualScope, Set%3CString%3E scopes) %7B%0A      if (scopes.isEmpty() %7C%7C scopes.contains(actualScope)) %7B%0A        return true;%0A      %7D%0A%0A      for (String scope : scopes) %7B%0A        if (actualScope.startsWith(scope)) %7B%0A          return true;%0A        %7D%0A      %7D%0A%0A      return false
+ing.matches(filter, metricResult.getKey());%0A        %7D%0A      %7D
 ;%0A  

Use MetricFiltering in Spark runner.