/*
 * 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
 *
 *   https://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.accumulo.core.dataImpl;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;
import static org.apache.accumulo.core.util.RowRangeUtil.requireKeyExtentDataRange;
import static org.apache.accumulo.core.util.RowRangeUtil.stripZeroTail;

import java.io.ByteArrayOutputStream;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.Comparator;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.UUID;

import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Range;
import org.apache.accumulo.core.data.TableId;
import org.apache.accumulo.core.data.TabletId;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.dataImpl.thrift.TKeyExtent;
import org.apache.accumulo.core.metadata.SystemTables;
import org.apache.accumulo.core.metadata.schema.MetadataSchema.TabletsSection;
import org.apache.accumulo.core.metadata.schema.MetadataSchema.TabletsSection.TabletColumnFamily;
import org.apache.accumulo.core.util.ByteBufferUtil;
import org.apache.accumulo.core.util.Pair;
import org.apache.accumulo.core.util.TextUtil;
import org.apache.hadoop.io.BinaryComparable;
import org.apache.hadoop.io.DataInputBuffer;
import org.apache.hadoop.io.DataOutputBuffer;
import org.apache.hadoop.io.Text;

/**
 * keeps track of information needed to identify a tablet
 */
public class KeyExtent implements Comparable<KeyExtent> {

  private static final String OBSCURING_HASH_ALGORITHM = "SHA-256";

  private final TableId tableId;
  private final Text endRow;
  private final Text prevEndRow;

  // lazily computed, as needed
  private volatile int hashCode = 0;

  private static final Text EMPTY_TEXT = new Text("");

  // The last tablet in a table has no end row, so null sorts last for end row; similarly, the first
  // tablet has no previous end row, so null sorts first for previous end row
  private static final Comparator<KeyExtent> COMPARATOR = Comparator.comparing(KeyExtent::tableId)
      .thenComparing(KeyExtent::endRow, Comparator.nullsLast(Text::compareTo))
      .thenComparing(KeyExtent::prevEndRow, Comparator.nullsFirst(Text::compareTo));

  /**
   * Create a new KeyExtent from its components.
   *
   * @param table the ID for the table
   * @param endRow the last row in this tablet, or null if this is the last tablet in this table
   * @param prevEndRow the last row in the immediately preceding tablet for the table, or null if
   *        this represents the first tablet in this table
   */
  public KeyExtent(TableId table, Text endRow, Text prevEndRow) {
    tableId = requireNonNull(table, "null table ID not allowed");
    if (endRow != null && prevEndRow != null && prevEndRow.compareTo(endRow) >= 0) {
      throw new IllegalArgumentException(
          "prevEndRow (" + prevEndRow + ") >= endRow (" + endRow + ")");
    }
    this.endRow = endRow == null ? null : new Text(endRow);
    this.prevEndRow = prevEndRow == null ? null : new Text(prevEndRow);
  }

  /**
   * Create a copy of a provided KeyExtent.
   *
   * @param original the KeyExtent to copy
   */
  public static KeyExtent copyOf(KeyExtent original) {
    return new KeyExtent(original.tableId(), original.endRow(), original.prevEndRow());
  }

  /**
   * Create a KeyExtent from its Thrift form.
   *
   * @param tke the KeyExtent in its Thrift object form
   */
  public static KeyExtent fromThrift(TKeyExtent tke) {
    TableId tableId = TableId.of(new String(ByteBufferUtil.toBytes(tke.table), UTF_8));
    Text endRow = tke.endRow == null ? null : new Text(ByteBufferUtil.toBytes(tke.endRow));
    Text prevEndRow =
        tke.prevEndRow == null ? null : new Text(ByteBufferUtil.toBytes(tke.prevEndRow));
    return new KeyExtent(tableId, endRow, prevEndRow);
  }

  /**
   * Convert to Thrift form.
   */
  public TKeyExtent toThrift() {
    return new TKeyExtent(ByteBuffer.wrap(tableId().canonical().getBytes(UTF_8)),
        endRow() == null ? null : TextUtil.getByteBuffer(endRow()),
        prevEndRow() == null ? null : TextUtil.getByteBuffer(prevEndRow()));
  }

  /**
   * Create a KeyExtent from a metadata previous end row entry.
   *
   * @param prevRowEntry a provided previous end row entry from the metadata table, stored in the
   *        <code>{@value TabletColumnFamily#STR_NAME}</code> column family and
   *        <code>{@value TabletColumnFamily#PREV_ROW_QUAL}</code> column. If the individual
   *        components are needed, consider {@link #KeyExtent(TableId, Text, Text)} or
   *        {@link #fromMetaRow(Text, Text)} instead.
   */
  public static KeyExtent fromMetaPrevRow(Entry<Key,Value> prevRowEntry) {
    return fromMetaRow(prevRowEntry.getKey().getRow(),
        TabletColumnFamily.decodePrevEndRow(prevRowEntry.getValue()));
  }

  /**
   * Create a KeyExtent from the table ID and the end row encoded in the row field of a tablet's
   * metadata entry, with no previous end row.
   *
   * @param encodedMetadataRow the encoded <code>tableId</code> and <code>endRow</code> from a
   *        metadata entry, as in <code>entry.getKey().getRow()</code> or from
   *        {@link TabletsSection#encodeRow(TableId, Text)}
   */
  public static KeyExtent fromMetaRow(Text encodedMetadataRow) {
    return fromMetaRow(encodedMetadataRow, (Text) null);
  }

  /**
   * Create a KeyExtent from the table ID and the end row encoded in the row field of a tablet's
   * metadata entry, along with a previous end row.
   *
   * @param encodedMetadataRow the encoded <code>tableId</code> and <code>endRow</code> from a
   *        metadata entry, as in <code>entry.getKey().getRow()</code> or from
   *        {@link TabletsSection#encodeRow(TableId, Text)}
   * @param prevEndRow the unencoded previous end row (a copy will be made)
   */
  public static KeyExtent fromMetaRow(Text encodedMetadataRow, Text prevEndRow) {
    Pair<TableId,Text> tableIdAndEndRow = TabletsSection.decodeRow(encodedMetadataRow);
    TableId tableId = tableIdAndEndRow.getFirst();
    Text endRow = tableIdAndEndRow.getSecond();
    return new KeyExtent(tableId, endRow, prevEndRow);
  }

  /**
   * Create a KeyExtent from a {@link TabletId}.
   *
   * @param tabletId the {@link TabletId} to convert to a KeyExtent
   */
  public static KeyExtent fromTabletId(TabletId tabletId) {
    if (tabletId instanceof TabletIdImpl) {
      return ((TabletIdImpl) tabletId).toKeyExtent();
    } else {
      return new KeyExtent(tabletId.getTable(), tabletId.getEndRow(), tabletId.getPrevEndRow());
    }
  }

  /**
   * Return a serialized form of the table ID and end row for this extent, in a form suitable for
   * use in the row portion of metadata entries for the tablet this extent represents.
   */
  public Text toMetaRow() {
    return TabletsSection.encodeRow(tableId(), endRow());
  }

  /**
   * Return the extent's table ID.
   */
  public TableId tableId() {
    return tableId;
  }

  /**
   * Returns this extent's end row
   */
  public Text endRow() {
    return endRow;
  }

  /**
   * Return the previous extent's end row
   */
  public Text prevEndRow() {
    return prevEndRow;
  }

  /**
   * Create a KeyExtent from a serialized form.
   *
   * @see #writeTo(DataOutput)
   */
  public static KeyExtent readFrom(DataInput in) throws IOException {
    Text tid = new Text();
    tid.readFields(in);
    TableId tableId = TableId.of(tid.toString());
    Text endRow = null;
    Text prevEndRow = null;
    boolean hasRow = in.readBoolean();
    if (hasRow) {
      endRow = new Text();
      endRow.readFields(in);
    }
    boolean hasPrevRow = in.readBoolean();
    if (hasPrevRow) {
      prevEndRow = new Text();
      prevEndRow.readFields(in);
    }
    return new KeyExtent(tableId, endRow, prevEndRow);
  }

  /**
   * Serialize this KeyExtent.
   *
   * @see #readFrom(DataInput)
   */
  public void writeTo(DataOutput out) throws IOException {
    new Text(tableId().canonical()).write(out);
    if (endRow() != null) {
      out.writeBoolean(true);
      endRow().write(out);
    } else {
      out.writeBoolean(false);
    }
    if (prevEndRow() != null) {
      out.writeBoolean(true);
      prevEndRow().write(out);
    } else {
      out.writeBoolean(false);
    }
  }

  @Override
  public int compareTo(KeyExtent other) {
    return COMPARATOR.compare(this, other);
  }

  @Override
  public int hashCode() {
    // lazily compute hash code, if needed
    if (hashCode != 0) {
      return hashCode;
    }
    int tmpHashCode = tableId().hashCode();
    tmpHashCode += endRow() == null ? 0 : endRow().hashCode();
    tmpHashCode += prevEndRow() == null ? 0 : prevEndRow().hashCode();
    hashCode = tmpHashCode;
    return hashCode;
  }

  @Override
  public boolean equals(Object o) {
    if (o == this) {
      return true;
    }
    if (!(o instanceof KeyExtent oke)) {
      return false;
    }
    return tableId().equals(oke.tableId()) && Objects.equals(endRow(), oke.endRow())
        && Objects.equals(prevEndRow(), oke.prevEndRow());
  }

  @Override
  public String toString() {
    String endRowString;
    String prevEndRowString;
    String tableIdString =
        tableId().canonical().replaceAll(";", "\\\\;").replaceAll("\\\\", "\\\\\\\\");

    if (endRow() == null) {
      endRowString = "<";
    } else {
      endRowString = ";" + TextUtil.truncate(endRow()).toString().replaceAll(";", "\\\\;")
          .replaceAll("\\\\", "\\\\\\\\");
    }

    if (prevEndRow() == null) {
      prevEndRowString = "<";
    } else {
      prevEndRowString = ";" + TextUtil.truncate(prevEndRow()).toString().replaceAll(";", "\\\\;")
          .replaceAll("\\\\", "\\\\\\\\");
    }

    return tableIdString + endRowString + prevEndRowString;
  }

  /**
   * Retrieve a unique identifier for this tablet that is useful for logging, without revealing the
   * contents of the end row and previous end row.
   */
  public UUID getUUID() {
    try (var baos = new ByteArrayOutputStream(); var dos = new DataOutputStream(baos)) {
      // to get a unique hash it is important to encode the data
      // like it is being serialized
      writeTo(dos);
      return UUID.nameUUIDFromBytes(baos.toByteArray());
    } catch (IOException impossible) {
      // impossible when ByteArrayOutputStream is backing the DataOutputStream
      throw new AssertionError(impossible);
    }
  }

  /**
   * Determine if another KeyExtent is wholly contained within the range of this one. * @param oke
   *
   * @return true if and only if the provided KeyExtent is contained within the range of this one
   */
  public boolean contains(KeyExtent oke) {
    // true if this prev row is before the other
    boolean containsPrevRow = prevEndRow() == null
        || (oke.prevEndRow() != null && prevEndRow().compareTo(oke.prevEndRow()) <= 0);
    // true if this end row is after the other
    boolean containsEndRow =
        endRow() == null || (oke.endRow() != null && endRow().compareTo(oke.endRow()) >= 0);
    return containsPrevRow && containsEndRow;
  }

  /**
   * Determine if the provided row is contained within the range of this KeyExtent.
   *
   * @return true if and only if the provided row is contained within the range of this one
   */
  public boolean contains(BinaryComparable row) {
    if (row == null) {
      throw new IllegalArgumentException(
          "Passing null to contains is ambiguous, could be in first or last extent of table");
    }
    return (prevEndRow() == null || prevEndRow().compareTo(row) < 0)
        && (endRow() == null || endRow().compareTo(row) >= 0);
  }

  /**
   * Compute a representation of this extent as a Range suitable for scanning this extent's table
   * and retrieving all data between this extent's previous end row (exclusive) and its end row
   * (inclusive) in that table (or another table in the same range, since the returned range does
   * not specify the table to scan).
   *
   * <p>
   * For example, if this extent represented a range of data from <code>A</code> to <code>Z</code>
   * for table <code>T</code>, the resulting Range could be used to scan table <code>T</code>'s data
   * in the exclusive-inclusive range <code>(A,Z]</code>, or the same range in another table,
   * <code>T2</code>.
   */
  public Range toDataRange() {
    return new Range(prevEndRow(), false, endRow(), true);
  }

  /**
   * Compute a representation of this extent as a Range suitable for scanning the corresponding
   * metadata table for this extent's table, and retrieving all the metadata for the table's tablets
   * between this extent's previous end row (exclusive) and its end row (inclusive) in that table.
   *
   * <p>
   * For example, if this extent represented a range of data from <code>A</code> to <code>Z</code>
   * for a user table, <code>T</code>, this would compute the range to scan
   * <code>accumulo.metadata</code> that would include all the metadata for <code>T</code>'s tablets
   * that contain data in the range <code>(A,Z]</code>.
   */
  public Range toMetaRange() {
    Text metadataPrevRow =
        TabletsSection.encodeRow(tableId(), prevEndRow() == null ? EMPTY_TEXT : prevEndRow());
    return new Range(metadataPrevRow, prevEndRow() == null, toMetaRow(), true);
  }

  /**
   * Creates a KeyExtent which represents the intersection of this KeyExtent and the passed in
   * range.
   * <p>
   * <b>Note:</b> The range provided must be a range that is derived from a KeyExtent. This means
   * the range must be in the format of a row range and also requires an exclusive start key, which
   * is the format that {@link #toDataRange()} uses
   *
   * @param range range to clip to
   * @return the intersection of this KeyExtent and the given range
   * @throws IllegalArgumentException if the KeyExtent and range do not overlap
   */
  public KeyExtent clip(Range range) {
    return clip(range, false);
  }

  /**
   * Creates a KeyExtent which represents the intersection of this KeyExtent and the passed in
   * range. Unlike {@link #clip(Range)}, this method can optionally return null if the given range
   * and this KeyExtent do not overlap, instead of throwing an exception. The returnNullIfDisjoint
   * parameter controls this behavior.
   * <p>
   * <b>Note:</b> The range provided must be a range that is derived from a KeyExtent. This means
   * the range must be in the format of a row range and also requires an exclusive start key, which
   * is the format that {@link #toDataRange()} uses
   *
   * @param range range to clip to
   * @param returnNullIfDisjoint true to return null if ranges are disjoint, false to throw an
   *        exception
   * @return the intersection of this KeyExtent and the given range, or null if given range and this
   *         KeyExtent do not overlap and returnNullIfDisjoint is true
   * @throws IllegalArgumentException if the KeyExtent and range does not overlap and
   *         returnNullIfDisjoint is false
   *
   * @see KeyExtent#clip(Range)
   **/
  public KeyExtent clip(Range range, boolean returnNullIfDisjoint) {
    // This will require a range that matches a row range generated by toDataRange()
    // This range itself will be required to be an inclusive start and exclusive end
    // The start and end rows will be required to be exclusive keys (ending in 0x00)
    requireKeyExtentDataRange(range);

    // If returnNullIfDisjoint is false then this will throw an exception if
    // the ranges are disjoint, otherwise we can just return null
    final Range clippedRange = this.toDataRange().clip(range, returnNullIfDisjoint);
    if (clippedRange == null) {
      return null;
    }

    // Build the new KeyExtent with the clipped range. We need to strip off the ending byte
    // which will essentially reverse what toDataRange() does
    Text endRow = clippedRange.getEndKey() != null
        ? new Text(stripZeroTail(clippedRange.getEndKey().getRowData()).toArray()) : null;
    Text prevEndRow = clippedRange.getStartKey() != null
        ? new Text(stripZeroTail(clippedRange.getStartKey().getRowData()).toArray()) : null;

    return new KeyExtent(tableId, endRow, prevEndRow);
  }

  private boolean startsAfter(KeyExtent other) {
    KeyExtent nke = requireNonNull(other);
    return tableId().compareTo(nke.tableId()) > 0 || (prevEndRow() != null && nke.endRow() != null
        && prevEndRow().compareTo(nke.endRow()) >= 0);
  }

  private static Text rowAfterPrevRow(KeyExtent nke) {
    Text row = new Text(nke.prevEndRow());
    row.append(new byte[] {0}, 0, 1);
    return row;
  }

  // Some duplication with TabletLocatorImpl
  public static Set<KeyExtent> findOverlapping(KeyExtent nke, SortedSet<KeyExtent> extents) {
    if (nke == null || extents == null || extents.isEmpty()) {
      return Collections.emptySet();
    }

    SortedSet<KeyExtent> start;

    if (nke.prevEndRow() != null) {
      Text row = rowAfterPrevRow(nke);
      KeyExtent lookupKey = new KeyExtent(nke.tableId(), row, null);
      start = extents.tailSet(lookupKey);
    } else {
      KeyExtent lookupKey = new KeyExtent(nke.tableId(), new Text(), null);
      start = extents.tailSet(lookupKey);
    }

    TreeSet<KeyExtent> result = new TreeSet<>();
    for (KeyExtent ke : start) {
      if (ke.startsAfter(nke)) {
        break;
      }
      result.add(ke);
    }
    return result;
  }

  public boolean overlaps(KeyExtent other) {
    SortedSet<KeyExtent> set = new TreeSet<>();
    set.add(other);
    return !findOverlapping(this, set).isEmpty();
  }

  // Specialization of findOverlapping(KeyExtent, SortedSet<KeyExtent> to work with SortedMap
  public static Set<KeyExtent> findOverlapping(KeyExtent nke, SortedMap<KeyExtent,?> extents) {
    if (nke == null || extents == null || extents.isEmpty()) {
      return Collections.emptySet();
    }

    SortedMap<KeyExtent,?> start;

    if (nke.prevEndRow() != null) {
      Text row = rowAfterPrevRow(nke);
      KeyExtent lookupKey = new KeyExtent(nke.tableId(), row, null);
      start = extents.tailMap(lookupKey);
    } else {
      KeyExtent lookupKey = new KeyExtent(nke.tableId(), new Text(), null);
      start = extents.tailMap(lookupKey);
    }

    TreeSet<KeyExtent> result = new TreeSet<>();
    for (Entry<KeyExtent,?> entry : start.entrySet()) {
      KeyExtent ke = entry.getKey();
      if (ke.startsAfter(nke)) {
        break;
      }
      result.add(ke);
    }
    return result;
  }

  public boolean isPreviousExtent(KeyExtent prevExtent) {
    if (prevExtent == null) {
      return prevEndRow() == null;
    }
    if (!prevExtent.tableId().equals(tableId())) {
      throw new IllegalArgumentException("Cannot compare across tables " + prevExtent + " " + this);
    }
    if (prevExtent.endRow() == null || prevEndRow() == null) {
      return false;
    }
    return prevExtent.endRow().equals(prevEndRow());
  }

  public boolean isSystemTable() {
    return SystemTables.containsTableId(tableId());
  }

  public boolean isMeta() {
    return tableId().equals(SystemTables.METADATA.tableId()) || isRootTablet();
  }

  public boolean isRootTablet() {
    return tableId().equals(SystemTables.ROOT.tableId());
  }

  public String obscured() {
    MessageDigest digester;
    try {
      digester = MessageDigest.getInstance(OBSCURING_HASH_ALGORITHM);
    } catch (NoSuchAlgorithmException e) {
      throw new IllegalStateException(e);
    }
    if (endRow() != null && endRow().getLength() > 0) {
      digester.update(endRow().getBytes(), 0, endRow().getLength());
    }
    return Base64.getEncoder().encodeToString(digester.digest());
  }

  public String toBase64() {
    DataOutputBuffer buffer = new DataOutputBuffer();
    try {
      writeTo(buffer);
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }

    return Base64.getEncoder().encodeToString(Arrays.copyOf(buffer.getData(), buffer.getLength()));
  }

  public static KeyExtent fromBase64(String encoded) {
    byte[] data = Base64.getDecoder().decode(encoded);
    DataInputBuffer buffer = new DataInputBuffer();
    buffer.reset(data, data.length);
    try {
      return KeyExtent.readFrom(buffer);
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }
  }

}
