#pragma once

#include "mlir/IR/Builders.h"

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/FunctionExtras.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>

namespace onnx_mlir {

class ConstantPool;

enum class StaticIntSequenceKind {
  Uniform,
  Affine,
  RunLengthEncoded,
  Dense
};

class StaticIntSequence {
public:
  static StaticIntSequence fromValues(llvm::ArrayRef<int64_t> values);
  static StaticIntSequence uniform(int64_t value, size_t count);
  static StaticIntSequence affine(int64_t base, int64_t step, size_t count);

  size_t size() const { return count; }
  int64_t valueAt(size_t index) const;
  std::optional<size_t> find(int64_t value, size_t begin, size_t length) const;
  StaticIntSequence slice(size_t begin, size_t count) const;
  StaticIntSequence remap(llvm::ArrayRef<unsigned> indices) const;
  StaticIntSequenceKind getKind() const { return kind; }

  bool operator==(const StaticIntSequence& other) const;
  llvm::hash_code hash() const;

  void forEachEqualRun(
    llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;

private:
  friend class StaticIntSequenceChain;
  friend void setStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef,
                                       const StaticIntSequence &, size_t);
  friend mlir::FailureOr<StaticIntSequence>
  getStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, size_t);

  static StaticIntSequence runLengthEncoded(
    llvm::ArrayRef<int64_t> runs, size_t count);
  StaticIntSequenceKind kind = StaticIntSequenceKind::Dense;
  size_t count = 0;
  int64_t base = 0;
  int64_t step = 0;
  llvm::SmallVector<int64_t> data;
};

struct StaticIntSequenceSlice {
  const StaticIntSequence *sequence = nullptr;
  size_t begin = 0;
  size_t count = 0;
};

class StaticIntSequenceChain {
public:
  void append(const StaticIntSequence &sequence, size_t begin, size_t count);
  void append(StaticIntSequence sequence);
  size_t size() const { return count; }
  int64_t valueAt(size_t index) const;
  void forEachSegment(llvm::function_ref<void(
    const StaticIntSequence &, size_t, size_t)> callback) const;
  void forEachEqualRun(
    llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
  StaticIntSequence canonicalize() const;

private:
  friend class StaticIntSequenceChainCursor;
  llvm::SmallVector<StaticIntSequenceSlice> slices;
  llvm::SmallVector<std::unique_ptr<StaticIntSequence>> owned;
  size_t count = 0;
};

class StaticIntSequenceChainCursor {
public:
  explicit StaticIntSequenceChainCursor(const StaticIntSequenceChain &chain)
    : chain(chain) {}

  bool done() const { return slice == chain.slices.size(); }
  int64_t value() const;
  void advance();

private:
  const StaticIntSequenceChain &chain;
  size_t slice = 0;
  size_t offset = 0;
};

void setStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
                              const StaticIntSequence &sequence,
                              size_t logicalCount);

mlir::FailureOr<StaticIntSequence>
getStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
                         size_t logicalCount);

mlir::Value emitStaticIntLookup(const StaticIntSequence& sequence,
                                mlir::Value position,
                                mlir::Operation* constantAnchor,
                                ConstantPool& constants,
                                mlir::OpBuilder& builder,
                                mlir::Location loc);

} // namespace onnx_mlir