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

package io.airbyte.integrations.destination.keen;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.joestelmach.natty.Parser;
import io.airbyte.protocol.models.AirbyteRecordMessage;
import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This class is used for timestamp inference. Keen leverages use of time-related data for it's
 * analytics, so it's important to have timestamp values for historical data if possible. If stream
 * contains cursor field, then its value is used as a timestamp, if parsing it is possible.
 */
public class KeenTimestampService {

  private static final Logger LOGGER = LoggerFactory.getLogger(KeenRecordsConsumer.class);

  private static final long SECONDS_FROM_EPOCH_THRESHOLD = 1_000_000_000L;

  private static final long MILLIS_FROM_EPOCH_THRESHOLD = 10_000_000_000L;

  // Map containing stream names paired with their cursor fields
  private Map<String, List<String>> streamCursorFields;
  private final Parser parser;
  private final boolean timestampInferenceEnabled;

  public KeenTimestampService(final ConfiguredAirbyteCatalog catalog, final boolean timestampInferenceEnabled) {
    this.streamCursorFields = new HashMap<>();
    this.parser = new Parser();
    this.timestampInferenceEnabled = timestampInferenceEnabled;

    if (timestampInferenceEnabled) {
      LOGGER.info("Initializing KeenTimestampService, finding cursor fields.");
      streamCursorFields = catalog.getStreams()
          .stream()
          .filter(stream -> stream.getCursorField().size() > 0)
          .map(s -> Pair.of(s.getStream().getName(), s.getCursorField()))
          .collect(Collectors.toMap(Pair::getKey, Pair::getValue));
    }
  }

  /**
   * Tries to inject keen.timestamp field to the given message data. If the stream contains cursor
   * field, it's value is tried to be parsed to timestamp. If this procedure fails, stream is removed
   * from timestamp-parsable stream map, so parsing is not tried for future messages in the same
   * stream. If parsing succeeds, keen.timestamp field is put as a JSON node to the message data and
   * whole data is returned. Otherwise, keen.timestamp is set to emittedAt value
   *
   * @param message AirbyteRecordMessage containing record data
   * @return Record data together with keen.timestamp field
   */
  public JsonNode injectTimestamp(final AirbyteRecordMessage message) {
    final String streamName = message.getStream();
    final List<String> cursorField = streamCursorFields.get(streamName);
    final JsonNode data = message.getData();
    if (timestampInferenceEnabled && cursorField != null) {
      try {
        final String timestamp = parseTimestamp(cursorField, data);
        injectTimestamp(data, timestamp);
      } catch (final Exception e) {
        // If parsing of timestamp has failed, remove stream from timestamp-parsable stream map,
        // so it won't be parsed for future messages.
        LOGGER.info("Unable to parse cursor field: {} into a keen.timestamp", cursorField);
        streamCursorFields.remove(streamName);
        injectTimestamp(data, Instant.ofEpochMilli(message.getEmittedAt()).toString());
      }
    } else {
      injectTimestamp(data, Instant.ofEpochMilli(message.getEmittedAt()).toString());
    }
    return data;
  }

  private void injectTimestamp(final JsonNode data, final String timestamp) {
    final ObjectNode root = ((ObjectNode) data);
    root.set("keen", JsonNodeFactory.instance.objectNode().put("timestamp", timestamp));
  }

  private String parseTimestamp(final List<String> cursorField, final JsonNode data) {
    final JsonNode timestamp = getNestedNode(data, cursorField);
    final long numberTimestamp = timestamp.asLong();
    // if cursor value is below given threshold, assume that it's not epoch timestamp but ordered id
    if (numberTimestamp >= SECONDS_FROM_EPOCH_THRESHOLD) {
      return dateFromNumber(numberTimestamp);
    }
    // if timestamp is 0, then parsing it to long failed - let's try with String now
    if (numberTimestamp == 0) {
      return parser
          .parse(timestamp.asText())
          .get(0).getDates()
          .get(0)
          .toInstant()
          .toString();
    }
    throw new IllegalStateException();
  }

  private String dateFromNumber(final Long timestamp) {
    // if cursor value is above given threshold, then assume that it's Unix timestamp in milliseconds
    if (timestamp > MILLIS_FROM_EPOCH_THRESHOLD) {
      return Instant.ofEpochMilli(timestamp).toString();
    }
    return Instant.ofEpochSecond(timestamp).toString();
  }

  private static JsonNode getNestedNode(final JsonNode data, final List<String> fieldNames) {
    return fieldNames.stream().reduce(data, JsonNode::get, (first, second) -> second);
  }

  public Map<String, List<String>> getStreamCursorFields() {
    return streamCursorFields;
  }

}
