// This file is part of OpenTSDB.
// Copyright (C) 2010-2018  The OpenTSDB Authors.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 2.1 of the License, or (at your
// option) any later version.  This program is distributed in the hope that it
// will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
// of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser
// General Public License for more details.  You should have received a copy
// of the GNU Lesser General Public License along with this program.  If not,
// see <http://www.gnu.org/licenses/>.

package net.opentsdb.uid;

import java.util.ArrayList;
import java.util.Arrays;

import com.stumbleupon.async.Deferred;
import net.opentsdb.core.TSDB;
import net.opentsdb.utils.Config;
import org.hbase.async.KeyValue;
import org.hbase.async.Scanner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Manages preloading of the UID cache.
 */
public class UidCachePreloader {

  private static final Logger LOG = LoggerFactory.getLogger(UidCachePreloader.class);

  // Prevent instantiation of this utility class:
  private UidCachePreloader() {
  }

  public static void preload(TSDB tsdb, Config config) {
    int maxRowsPerBatch = config.get_preload_uid_cache_max_rows_per_batch();
    long batchTimeoutMillis = config.get_preload_uid_cache_per_batch_timeout_ms();

    LOG.info("Preload UID cache start [Batch size=" + maxRowsPerBatch + ", per-batch-timeout="
        + batchTimeoutMillis + " ms.]");
    long startTime = System.currentTimeMillis();

    long metricsCount = preload(new ScannerFactory(tsdb.getMetrics(), maxRowsPerBatch),
        batchTimeoutMillis, "metrics");
    long tagNamesCount = preload(new ScannerFactory(tsdb.getTagNames(), maxRowsPerBatch),
        batchTimeoutMillis, "tag name");
    long tagValuesCount = preload(new ScannerFactory(tsdb.getTagValues(), maxRowsPerBatch),
        batchTimeoutMillis, "tag value");

    long stopTime = System.currentTimeMillis();
    double elapsedTimeSecs = (double)(stopTime - startTime) / 1000.0;
    LOG.info(String.format("Preload UID cache done: metrics=%d, tag names=%d, tag values=%d",
        metricsCount, tagNamesCount, tagValuesCount));
    LOG.info(String.format("Preload UID cache total duration=%.2f secs.", elapsedTimeSecs));
  }

  private static long preload(ScannerFactory scannerFactory, long batchTimeoutMillis,
      String uidType) {
    LOG.info("Start preloading " + uidType + " UIDs ...");
    long startTime = System.currentTimeMillis();
    long uidCount = preload(scannerFactory, batchTimeoutMillis);
    long stopTime = System.currentTimeMillis();
    double elapsedTimeSecs = (double)(stopTime - startTime) / 1000.0;
    LOG.info(String.format("Preloading %s UIDs done; count=%d; duration=%.2f secs.",
        uidType, uidCount, elapsedTimeSecs));
    return uidCount;
  }

  private static long preload(ScannerFactory scannerFactory, long batchTimeoutMillis) {
    Scanner scanner = scannerFactory.get();
    String kind = scannerFactory.getUid().kind();
    if (LOG.isDebugEnabled()) {
      LOG.debug("Starting to retrieve rows for UID kind " + kind);
    }
    long totalUidCount = 0;
    boolean continueScanning = true;
    while (continueScanning) {
      Deferred<ArrayList<ArrayList<KeyValue>>> deferredScannedRows = scanner.nextRows();
      try {
        ArrayList<ArrayList<KeyValue>> scannedRows =
            deferredScannedRows.joinUninterruptibly(batchTimeoutMillis);
        if (null == scannedRows) {
          // No more rows to scan: we are done:
          continueScanning = false;
          if (LOG.isDebugEnabled()) {
            LOG.debug("Done: no more rows for UID kind " + kind);
          }
        } else {
          cacheUids(scannedRows, scannerFactory);
          long currentUidCount = scannedRows.size();
          totalUidCount += currentUidCount;
          if (LOG.isDebugEnabled()) {
            LOG.debug("Retrieved and cached next round of " + currentUidCount
                + " rows for UID kind " + kind);
          }
        }
      } catch (com.stumbleupon.async.TimeoutException e) {
        // Timed out:
        LOG.warn("Preload of an incremental round of " + kind + " UIDs timed out with message \""
            + e.getMessage() + "\"; bailing out ...");
        continueScanning = false;
        // TODO(ankan): Implement retry with exponential backoff instead of bailing out.
      } catch (Throwable e) {
        // This will happen if the deferred result of scanner.nextRows() is an error, or
        // because of any other runtime exceptions from the above steps.
        LOG.error("Preload of an incremental round of " + kind + " UIDs failed with error \""
            + e.getMessage() + "\"; bailing out ...");
        continueScanning = false;
      }
    }
    scanner.close();
    // NOTE(ankan): The scanner.close() returns a Deferred, and it's probably better protocol
    // to wait for the result (success/failure) to become available before proceeding further;
    // however, I've chosen not to do so, since the current version of AsyncHBase waits for the
    // RegionServer to close the scanner before it completes the wait, and that might add
    // significantly to the latency.  If this is seen to cause any problems, we should revisit
    // this choice.
    return totalUidCount;
  }

  private static void cacheUids(ArrayList<ArrayList<KeyValue>> uidRows,
      ScannerFactory scannerFactory) {
    UniqueId uid = scannerFactory.getUid();
    Scanner scanner = scannerFactory.get();

    for (ArrayList<KeyValue> row : uidRows) {
      if (1 != row.size()) {
        LOG.error("Scanner " + scanner.toString()
            + " returned a row that doesn't have exactly one KeyValue: " + row);
        if (row.isEmpty()) {
          continue;
        }
      }
      byte[] key = row.get(0).key();
      String name = UniqueId.fromBytes(key);
      byte[] id = row.get(0).value();
      byte[] cachedId = uid.getIdFromCache(name);
      if (null == cachedId) {
        uid.cacheMapping(name, id); // The main work is done here
      } else if (!Arrays.equals(id, cachedId)) {
        throw new IllegalStateException("For kind=" + uid.kind()
            + " name=" + name + ", we have id=" + Arrays.toString(cachedId)
            + " in cache, but just scanned id=" + Arrays.toString(id));
      }
    }
  }

  /**
   * Creates a new {@link Scanner} for the given uid and maxResults.
   *
   * This should be created afresh for every new scan, and not reused, since the scanner
   * it holds cannot be reused.
   */
  private static class ScannerFactory {

    private final Scanner scanner;
    private final UniqueId uid;

    ScannerFactory(UniqueId uid, int maxRowsPerBatch) {
      this.uid = uid;
      Scanner scanner = uid.getUidCachePreloadScanner(maxRowsPerBatch);
      if (LOG.isDebugEnabled()) {
        LOG.debug("Created new scanner " + scanner.toString() + " for " + uid.kind()
            + " UID cache preloading, in batches of " + String.valueOf(maxRowsPerBatch) + " UIDs");
      }
      this.scanner = scanner;
    }

    Scanner get() {
      return scanner;
    }

    UniqueId getUid() {
      return uid;
    }
  }
}
