/*
 * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 */

package io.airbyte.commons.logging;

import io.airbyte.commons.logging.LoggingHelper.Color;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import org.slf4j.MDC;

/**
 * This class is an autoClosable class that will add some specific values into the log MDC. When
 * being close, it will restore the orginal MDC. It is advise to use it like that:
 *
 * <pre>
 *   <code>
 *     try(final ScopedMDCChange scopedMDCChange = new ScopedMDCChange(
 *      new HashMap&lt;String, String&gt;() {{
 *        put("my", "value");
 *      }}
 *     )) {
 *        ...
 *     }
 *   </code>
 * </pre>
 */
public class MdcScope implements AutoCloseable {

  public final static MdcScope.Builder DEFAULT_BUILDER = new Builder();

  private final Map<String, String> originalContextMap;

  public MdcScope(final Map<String, String> keyValuesToAdd) {
    originalContextMap = MDC.getCopyOfContextMap();

    keyValuesToAdd.forEach(
        (key, value) -> MDC.put(key, value));
  }

  @Override
  public void close() {
    MDC.setContextMap(originalContextMap);
  }

  public static class Builder {

    private Optional<String> maybeLogPrefix = Optional.empty();
    private Optional<Color> maybePrefixColor = Optional.empty();
    private boolean simple = true;

    public Builder setLogPrefix(final String logPrefix) {
      this.maybeLogPrefix = Optional.ofNullable(logPrefix);

      return this;
    }

    public Builder setPrefixColor(final Color color) {
      this.maybePrefixColor = Optional.ofNullable(color);

      return this;
    }

    // Use this to disable simple logging for things in an MdcScope.
    // If you're using this, you're probably starting to use MdcScope outside of container labelling.
    // If so, consider changing the defaults / builder / naming.
    public Builder setSimple(final boolean simple) {
      this.simple = simple;

      return this;
    }

    public MdcScope build() {
      final Map<String, String> extraMdcEntries = new HashMap<>();

      maybeLogPrefix.ifPresent(logPrefix -> {
        final String potentiallyColoredLog = maybePrefixColor
            .map(color -> LoggingHelper.applyColor(color, logPrefix))
            .orElse(logPrefix);

        extraMdcEntries.put(LoggingHelper.LOG_SOURCE_MDC_KEY, potentiallyColoredLog);

        if (simple) {
          // outputs much less information for this line. see log4j2.xml to see exactly what this does
          extraMdcEntries.put("simple", "true");
        }
      });

      return new MdcScope(extraMdcEntries);
    }

  }

}
