//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.

//  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 <cassert>

#include <cstdint>

#include <unordered_map>

#include "db/blob/blob_constants.h"

#include "rocksdb/rocksdb_namespace.h"

#include "rocksdb/status.h"

namespace ROCKSDB_NAMESPACE {

class BlobIndex;

struct ParsedInternalKey;

class Slice;

// A class that can be used to compute the amount of additional garbage

// generated by a compaction. It parses the keys and blob references in the

// input and output of a compaction, and aggregates the "inflow" and "outflow"

// on a per-blob file basis. The amount of additional garbage for any given blob

// file can then be computed by subtracting the outflow from the inflow.

class BlobGarbageMeter {

 public:

  // A class to store the number and total size of blobs on a per-blob file

  // basis.

  class BlobStats {

   public:

    void Add(uint64_t bytes) {

      ++count_;

      bytes_ += bytes;

    }

    void Add(uint64_t count, uint64_t bytes) {

      count_ += count;

      bytes_ += bytes;

    }

    uint64_t GetCount() const { return count_; }

    uint64_t GetBytes() const { return bytes_; }

   private:

    uint64_t count_ = 0;

    uint64_t bytes_ = 0;

  };

  // A class to keep track of the "inflow" and the "outflow" and to compute the

  // amount of additional garbage for a given blob file.

  class BlobInOutFlow {

   public:

    void AddInFlow(uint64_t bytes) {

      in_flow_.Add(bytes);

      assert(IsValid());

    }

    void AddOutFlow(uint64_t bytes) {

      out_flow_.Add(bytes);

      assert(IsValid());

    }

    const BlobStats& GetInFlow() const { return in_flow_; }

    const BlobStats& GetOutFlow() const { return out_flow_; }

    bool IsValid() const {

      return in_flow_.GetCount() >= out_flow_.GetCount() &&

             in_flow_.GetBytes() >= out_flow_.GetBytes();

    }

    bool HasGarbage() const {

      assert(IsValid());

      return in_flow_.GetCount() > out_flow_.GetCount();

    }

    uint64_t GetGarbageCount() const {

      assert(IsValid());

      assert(HasGarbage());

      return in_flow_.GetCount() - out_flow_.GetCount();

    }

    uint64_t GetGarbageBytes() const {

      assert(IsValid());

      assert(HasGarbage());

      return in_flow_.GetBytes() - out_flow_.GetBytes();

    }

   private:

    BlobStats in_flow_;

    BlobStats out_flow_;

  };

  Status ProcessInFlow(const Slice& key, const Slice& value);

  Status ProcessOutFlow(const Slice& key, const Slice& value);

  const std::unordered_map<uint64_t, BlobInOutFlow>& flows() const {

    return flows_;

  }

 private:

  static Status GetBlobReferenceDetails(const ParsedInternalKey& ikey,

                                        const BlobIndex& blob_index,

                                        uint64_t* blob_file_number,

                                        uint64_t* bytes);

  static Status ParseBlobIndexReference(const ParsedInternalKey& ikey,

                                        const Slice& value,

                                        uint64_t* blob_file_number,

                                        uint64_t* bytes);

  void AddFlow(uint64_t blob_file_number, uint64_t bytes, bool is_inflow);

  Status ProcessFlow(const Slice& key, const Slice& value, bool is_inflow);

  Status ProcessEntityBlobReferences(const ParsedInternalKey& ikey,

                                     const Slice& value, bool is_inflow);

  std::unordered_map<uint64_t, BlobInOutFlow> flows_;

};

}  // namespace ROCKSDB_NAMESPACE