/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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.
 */
#pragma once

#include "velox/common/base/Exceptions.h"
#include "velox/common/base/SimdUtil.h"
#include "velox/vector/FlatVector.h"
#include "velox/vector/SimpleVector.h"
#include "velox/vector/TypeAliases.h"

namespace facebook::velox {

/**
 * Provides a vector type that enacts run-length encoding for vectors that have
 * repeated runs of the same value.  For example, this would be useful for
 * {foo, foo, foo, foo, bar, bar, bar, foo, foo, foo, foo, foo} as it can be
 * reduced to {foo, bar, foo} along with run length information of {4, 3, 5}.
 *
 * NOTE - current implementation is not thread safe when accessing values.
 */
template <typename T>
class SequenceVector : public SimpleVector<T> {
 public:
  static constexpr bool can_simd =
      (std::is_same_v<T, int64_t> || std::is_same_v<T, int32_t> ||
       std::is_same_v<T, int16_t> || std::is_same_v<T, int8_t> ||
       std::is_same_v<T, size_t>);

  SequenceVector(
      velox::memory::MemoryPool* pool,
      size_t length,
      VectorPtr sequenceValues,
      BufferPtr sequenceLengths,
      const SimpleVectorStats<T>& stats = {},
      std::optional<int32_t> distinctCount = std::nullopt,
      std::optional<vector_size_t> nullCount = std::nullopt,
      std::optional<bool> sorted = std::nullopt,
      std::optional<ByteCount> representedBytes = std::nullopt,
      std::optional<ByteCount> storageByteCount = std::nullopt);

 public:
  ~SequenceVector() override = default;

  bool mayHaveNulls() const override {
    return sequenceValues_->mayHaveNulls();
  }

  bool mayHaveNullsRecursive() const override {
    return sequenceValues_->mayHaveNullsRecursive();
  }

  bool isNullAtFast(vector_size_t idx) const;

  bool isNullAt(vector_size_t idx) const override {
    return isNullAtFast(idx);
  }

  bool containsNullAt(vector_size_t idx) const override {
    if constexpr (std::is_same_v<T, ComplexType>) {
      if (isNullAt(idx)) {
        return true;
      }

      size_t offset = offsetOfIndex(idx);
      return sequenceValues_->containsNullAt(offset);
    } else {
      return isNullAt(idx);
    }
  }

  const T valueAtFast(vector_size_t idx) const;

  const T valueAt(vector_size_t idx) const override {
    return valueAtFast(idx);
  }

  std::unique_ptr<SimpleVector<uint64_t>> hashAll() const override;

  /**
   * Loads a 256bit vector of data at the virtual byteOffset given
   * Note this method is implemented on each vector type, but is intentionally
   * not virtual for performance reasons
   *
   * @param byteOffset - the byte offset to laod from
   */
  xsimd::batch<T> loadSIMDValueBufferAt(size_t index) const;

  /**
   * Returns a shared_ptr to the underlying byte buffer holding the values for
   * this vector as SequenceLength types. This is used during execution to
   * process over the subset of values when possible.
   */
  inline VectorPtr getSequenceValues() const {
    return sequenceValues_;
  }

  BaseVector* loadedVector() override {
    auto loaded = BaseVector::loadedVectorShared(sequenceValues_);
    if (loaded == sequenceValues_) {
      return this;
    }
    sequenceValues_ = loaded;
    setInternalState();
    return this;
  }

  const BaseVector* loadedVector() const override {
    return const_cast<SequenceVector<T>*>(this)->loadedVector();
  }

  const VectorPtr& valueVector() const override {
    return sequenceValues_;
  }

  BufferPtr getSequenceLengths() const {
    return sequenceLengths_;
  }

  /**
   * Returns a shared_ptr to the underlying values container.
   */
  inline VectorPtr getValuesVectorShared() const {
    return sequenceValues_;
  }

  vector_size_t numSequences() const {
    return sequenceValues_->size();
  }

  const BufferPtr& wrapInfo() const override {
    return sequenceLengths_;
  }

  uint64_t retainedSize() const override {
    return sequenceValues_->retainedSize() + sequenceLengths_->capacity();
  }

  bool isScalar() const override {
    return sequenceValues_->isScalar();
  }

  /**
   * @returns the physical index for the location of the sequence value that
   * is logically resident at the given index, or, -1 for an index that is < 0
   * or an index that is >= the logical length of this vector.
   */
  vector_size_t offsetOfIndex(vector_size_t idx) const;

  const BaseVector* wrappedVector() const override {
    return sequenceValues_->wrappedVector();
  }

  vector_size_t wrappedIndex(vector_size_t index) const override {
    return sequenceValues_->wrappedIndex(offsetOfIndex(index));
  }

  void addNulls(const uint64_t* bits, const SelectivityVector& rows) override {
    throw std::runtime_error("addNulls not supported");
  }

  void addNulls(const SelectivityVector& rows) override {
    throw std::runtime_error("addNulls not supported");
  }

  std::string toString(vector_size_t index) const override {
    if (BaseVector::isNullAt(index)) {
      return "null";
    }
    auto inner = offsetOfIndex(index);
    std::stringstream out;
    out << "[" << index << "->" << inner << "] "
        << sequenceValues_->toString(inner);
    return out.str();
  }

  VectorPtr slice(vector_size_t, vector_size_t) const override {
    VELOX_NYI();
  }

  bool isNullsWritable() const override {
    return false;
  }

  VectorPtr copyPreserveEncodings(
      velox::memory::MemoryPool* pool = nullptr) const override {
    auto selfPool = pool ? pool : BaseVector::pool_;
    return std::make_shared<SequenceVector<T>>(
        selfPool,
        BaseVector::length_,
        sequenceValues_->copyPreserveEncodings(),
        AlignedBuffer::copy(selfPool, sequenceLengths_),
        SimpleVector<T>::stats_,
        BaseVector::distinctValueCount_,
        BaseVector::nullCount_,
        SimpleVector<T>::isSorted_,
        BaseVector::representedByteCount_,
        BaseVector::storageByteCount_);
  }

 private:
  // Prepares for use after construction.
  void setInternalState();

  bool checkLoadRange(size_t idx, size_t count) const;

  VectorPtr sequenceValues_;
  SimpleVector<T>* scalarSequenceValues_ = nullptr;
  BufferPtr sequenceLengths_;
  // Caches 'sequenceLengths_->as<vector_size_t>()'
  const vector_size_t* lengths_ = nullptr;

  // manage the state of the offset index operations to allow for the common
  // path of iterating in order to not require repeated trips through the
  // sequenceLengths_ array
  mutable vector_size_t lastIndexRangeStart_ = 0;
  mutable vector_size_t lastIndexRangeEnd_ = 0;
  mutable vector_size_t lastRangeIndex_ = 0;
};

template <typename T>
using SequenceVectorPtr = std::shared_ptr<SequenceVector<T>>;

} // namespace facebook::velox

#include "velox/vector/SequenceVector-inl.h"
