// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#ifndef TOOLS_TRACKING_MEMORY_MANAGER_H_
#define TOOLS_TRACKING_MEMORY_MANAGER_H_

#include <jxl/memory_manager.h>

#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <mutex>
#include <unordered_map>

#include "lib/jxl/base/status.h"

namespace jpegxl {
namespace tools {

const uint64_t kGiB = 1u << 30;

class TrackingMemoryManager {
 public:
  explicit TrackingMemoryManager(uint64_t cap = 0, uint64_t total_cap = 0);

  // void setInner(JxlMemoryManager* inner) { inner_ = inner; }

  JxlMemoryManager* get() { return &outer_; }

  jxl::Status Reset();

  bool seen_oom = false;
  uint64_t max_bytes_in_use = 0;
  uint64_t total_allocations = 0;
  uint64_t total_bytes_allocated = 0;

 private:
  static void* Alloc(void* opaque, size_t size);
  static void Free(void* opaque, void* address);

  std::unordered_map<void*, size_t> allocations_;
  std::mutex numbers_mutex_;
  std::mutex map_mutex_;
  uint64_t cap_;
  uint64_t total_cap_;
  uint64_t bytes_in_use_ = 0;
  uint64_t num_allocations_ = 0;
  JxlMemoryManager outer_;
  JxlMemoryManager default_;
  JxlMemoryManager* inner_;
};

}  // namespace tools
}  // namespace jpegxl

#endif  // TOOLS_TRACKING_MEMORY_MANAGER_H_