// Copyright 2020 Google LLC.
//
// 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 com.google.cloud.healthcare.etl.pipeline;

import com.google.cloud.healthcare.etl.model.mapping.Mappable;
import com.google.cloud.healthcare.etl.model.mapping.MappedFhirMessageWithSourceTime;
import com.google.cloud.healthcare.etl.model.mapping.MappingOutput;
import com.google.cloud.healthcare.etl.provider.mapping.MappingConfigProvider;
import com.google.cloud.healthcare.etl.provider.mapping.MappingConfigProviderFactory;
import com.google.cloud.healthcare.etl.util.library.TransformWrapper;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;
import org.apache.beam.sdk.metrics.Distribution;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * The core function of the mapping pipeline. Input is expected to be a parsed message. At this
 * moment, only higher level language (whistle) is supported.
 */
public class MappingFn<M extends Mappable> extends ErrorEnabledDoFn<M, MappingOutput> {
  // TODO(b/173141038): refactor the class for capturing performance metrics better.
  private static final Logger LOGGER = LoggerFactory.getLogger(MappingFn.class);
  public static final TupleTag<MappingOutput> MAPPING_TAG = new TupleTag<>("mapping");
  private final Distribution transformMetrics = Metrics.distribution(MappingFn.class, "Transform");

  // Ensure the initialization only happens once. Ideally this should be handled by the library.
  private static final AtomicBoolean initializeStarted = new AtomicBoolean();
  private static final AtomicBoolean initializeFinished = new AtomicBoolean();
  private static final CountDownLatch initializeGate = new CountDownLatch(1);

  private final ValueProvider<String> mappingPath;
  private final ValueProvider<String> mappings;
  private final boolean enablePerformanceMetrics;

  private TransformWrapper engine;

  // The config parameter should be the string representation of the whole mapping config, including
  // harmonization and libraries.
  protected MappingFn(
      ValueProvider<String> mappingPath,
      ValueProvider<String> mappings,
      Boolean enablePerformanceMetrics) {
    this.mappingPath = mappingPath;
    this.mappings = mappings;
    this.enablePerformanceMetrics = enablePerformanceMetrics;
  }

  public static MappingFn of(ValueProvider<String> mappingPath, Boolean enablePerformanceMetrics) {
    return new MappingFn(mappingPath, StaticValueProvider.of(""), enablePerformanceMetrics);
  }

  public static MappingFn of(String mappingPath, Boolean enablePerformanceMetrics) {
    return of(StaticValueProvider.of(mappingPath), enablePerformanceMetrics);
  }

  public static MappingFn of(
      ValueProvider<String> mappingPath,
      ValueProvider<String> mappings,
      Boolean enablePerformanceMetrics) {
    return new MappingFn(mappingPath, mappings, enablePerformanceMetrics);
  }

  public static MappingFn of(
      String mappings, String mappingPath, Boolean enablePerformanceMetrics) {
    return of(
        StaticValueProvider.of(mappingPath),
        StaticValueProvider.of(mappings),
        enablePerformanceMetrics);
  }

  @Setup
  public void initialize() throws InterruptedException {
    // Make sure the mapping configuration is only initialized once.
    if (initializeStarted.compareAndSet(false, true)) {
      LOGGER.info("Initializing the mapping configurations.");
      engine = TransformWrapper.getInstance();
      try {
        // Mapping configurations are loaded from the `mappingPath` only if `mappings` is absent.
        String mappingsToUse = mappings.get();
        if (Strings.isNullOrEmpty(mappingsToUse)) {
          mappingsToUse = loadMapping(mappingPath.get());
        }
        engine.initializeWhistler(mappingsToUse);
        // This has to be done before opening the gate to avoid race conditions.
        initializeFinished.compareAndSet(false, true);
      } catch (RuntimeException e) {
        LOGGER.error("Unable to initialize mapping configurations.", e);
        // Reset the flag so that initialization can be retried (e.g. users fixed the mapping
        // configurations), by another thread.
        initializeStarted.compareAndSet(true, false);
        throw e; // Fail fast.
      } finally {
        // Open the gate in case other threads are waiting. Successive transformations will fail if
        // initialization failed.
        initializeGate.countDown();
      }
    } else if (!initializeFinished.get()) {
      // Wait for the gate to open if initialization is not done.
      LOGGER.info("Waiting for initialization to finish.");
      initializeGate.await();
    }
  }

  private static String loadMapping(String mappingPath) {
    MappingConfigProvider provider = MappingConfigProviderFactory.createProvider(mappingPath);
    try {
      return new String(provider.getMappingConfig(true /* force */), StandardCharsets.UTF_8);
    } catch (IOException e) {
      throw new RuntimeException("Unable to load mapping configurations.", e);
    }
  }

  // Runs a lambda and collects the metrics.
  private static <T> T runAndReportMetrics(Distribution metrics, Supplier<T> supplier) {
    Instant start = Instant.now();
    T result = supplier.get();
    metrics.update(Instant.now().toEpochMilli() - start.toEpochMilli());
    return result;
  }

  @Override
  public void process(ProcessContext ctx) {
    M input = ctx.element();
    ctx.output(
        runAndReportMetrics(
            transformMetrics,
            () -> {
              String transformedData = engine.transform(input.getData());
              return (enablePerformanceMetrics)
                  ? new MappedFhirMessageWithSourceTime(
                      transformedData, input.getCreateTime().get())
                  : new MappedFhirMessageWithSourceTime(transformedData);
            }));
  }

  @Override
  protected List<String> getSources(M input) {
    return Collections.singletonList(input.getId());
  }
}

@@ -2265,29 +2265,31 @@
 ping%22);%0A  pr
-ivate
+otected
  final Distr
@@ -2314,16 +2314,22 @@
 etrics =
+%0A     
  Metrics
@@ -2861,21 +2861,23 @@
 s;%0A%0A  pr
-ivate
+otected
  Transfo
@@ -6084,29 +6084,31 @@
 etrics.%0A  pr
-ivate
+otected
  static %3CT%3E 
@@ -6978,9 +6978,8 @@
 );%0A  %7D%0A%7D
-%0A

Change MappingFn fields to be protected instead of private.