//  Copyright (c) Meta Platforms, Inc. and affiliates.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).

#pragma once

#include "rocksdb/cache.h"
#include "rocksdb/secondary_cache.h"

namespace ROCKSDB_NAMESPACE {

// A SecondaryCache that implements stacking of a compressed secondary cache
// and a non-volatile (local flash) cache. It implements an admission
// policy of warming the bottommost tier (local flash) with compressed
// blocks from the SST on misses, and on hits in the bottommost tier,
// promoting to the compressed and/or primary block cache. The admission
// policies of the primary block cache and compressed secondary cache remain
// unchanged - promote on second access. There is no demotion ofablocks
// evicted from a tier. They are just discarded.
//
// In order to properly handle compressed blocks directly read from SSTs, and
// to allow writeback of blocks compressed by the compressed secondary
// cache in the future, we make use of the compression type and source
// cache tier arguments in InsertSaved.
class TieredSecondaryCache : public SecondaryCacheWrapper {
 public:
  TieredSecondaryCache(std::shared_ptr<SecondaryCache> comp_sec_cache,
                       std::shared_ptr<SecondaryCache> nvm_sec_cache,
                       TieredAdmissionPolicy adm_policy)
      : SecondaryCacheWrapper(comp_sec_cache), nvm_sec_cache_(nvm_sec_cache) {
#ifndef NDEBUG
    assert(adm_policy == TieredAdmissionPolicy::kAdmPolicyThreeQueue);
#else
    (void)adm_policy;
#endif
  }

  ~TieredSecondaryCache() override {}

  const char* Name() const override { return "TieredSecondaryCache"; }

  // This is a no-op as we currently don't allow demotion (i.e
  // insertion by the upper layer) of evicted blocks.
  virtual Status Insert(const Slice& /*key*/, Cache::ObjectPtr /*obj*/,
                        const Cache::CacheItemHelper* /*helper*/,
                        bool /*force_insert*/) override {
    return Status::OK();
  }

  // Warm up the nvm tier directly
  virtual Status InsertSaved(
      const Slice& key, const Slice& saved,
      CompressionType type = CompressionType::kNoCompression,
      CacheTier source = CacheTier::kVolatileTier) override {
    return nvm_sec_cache_->InsertSaved(key, saved, type, source);
  }

  virtual std::unique_ptr<SecondaryCacheResultHandle> Lookup(
      const Slice& key, const Cache::CacheItemHelper* helper,
      Cache::CreateContext* create_context, bool wait, bool advise_erase,
      Statistics* stats, bool& kept_in_sec_cache) override;

  virtual void WaitAll(
      std::vector<SecondaryCacheResultHandle*> handles) override;

 private:
  struct CreateContext : public Cache::CreateContext {
    const Slice* key;
    bool advise_erase;
    const Cache::CacheItemHelper* helper;
    Cache::CreateContext* inner_ctx;
    std::shared_ptr<SecondaryCacheResultHandle> inner_handle;
    SecondaryCache* comp_sec_cache;
    Statistics* stats;
  };

  class ResultHandle : public SecondaryCacheResultHandle {
   public:
    ~ResultHandle() override {}

    bool IsReady() override {
      if (inner_handle_ && inner_handle_->IsReady()) {
        Complete();
      }
      return ready_;
    }

    void Wait() override {
      inner_handle_->Wait();
      Complete();
    }

    size_t Size() override { return size_; }

    Cache::ObjectPtr Value() override { return value_; }

    void Complete() {
      size_ = inner_handle_->Size();
      value_ = inner_handle_->Value();
      inner_handle_.reset();
      ready_ = true;
    }

    void SetInnerHandle(std::unique_ptr<SecondaryCacheResultHandle>&& handle) {
      inner_handle_ = std::move(handle);
    }

    void SetSize(size_t size) { size_ = size; }

    void SetValue(Cache::ObjectPtr val) { value_ = val; }

    CreateContext* ctx() { return &ctx_; }

    SecondaryCacheResultHandle* inner_handle() { return inner_handle_.get(); }

   private:
    std::unique_ptr<SecondaryCacheResultHandle> inner_handle_;
    CreateContext ctx_;
    size_t size_;
    Cache::ObjectPtr value_;
    bool ready_ = false;
  };

  static void NoopDelete(Cache::ObjectPtr /*obj*/,
                         MemoryAllocator* /*allocator*/) {
    assert(false);
  }
  static size_t ZeroSize(Cache::ObjectPtr /*obj*/) {
    assert(false);
    return 0;
  }
  static Status NoopSaveTo(Cache::ObjectPtr /*from_obj*/,
                           size_t /*from_offset*/, size_t /*length*/,
                           char* /*out_buf*/) {
    assert(false);
    return Status::OK();
  }
  static Status MaybeInsertAndCreate(const Slice& data, CompressionType type,
                                     CacheTier source,
                                     Cache::CreateContext* ctx,
                                     MemoryAllocator* allocator,
                                     Cache::ObjectPtr* out_obj,
                                     size_t* out_charge);

  static const Cache::CacheItemHelper* GetHelper() {
    const static Cache::CacheItemHelper basic_helper(CacheEntryRole::kMisc,
                                                     &NoopDelete);
    const static Cache::CacheItemHelper maybe_insert_and_create_helper{
        CacheEntryRole::kMisc, &NoopDelete,           &ZeroSize,
        &NoopSaveTo,           &MaybeInsertAndCreate, &basic_helper,
    };
    return &maybe_insert_and_create_helper;
  }

  std::shared_ptr<SecondaryCache> comp_sec_cache_;
  std::shared_ptr<SecondaryCache> nvm_sec_cache_;
};

}  // namespace ROCKSDB_NAMESPACE