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

#include "tools/tracking_memory_manager.h"

#include <jxl/memory_manager.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <mutex>

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

namespace jpegxl {
namespace tools {

TrackingMemoryManager::TrackingMemoryManager(uint64_t cap, uint64_t total_cap)
    : cap_(cap), total_cap_(total_cap) {
  jxl::Status status = jxl::MemoryManagerInit(&default_, nullptr);
  JXL_DASSERT(status);
  (void)status;
  inner_ = &default_;

  outer_.opaque = reinterpret_cast<void*>(this);
  outer_.alloc = &Alloc;
  outer_.free = &Free;
}

void* TrackingMemoryManager::Alloc(void* opaque, size_t size) {
  if (opaque == nullptr) {
    JXL_DEBUG_ABORT("Internal logic error");
    return nullptr;
  }
  TrackingMemoryManager* self =
      reinterpret_cast<TrackingMemoryManager*>(opaque);
  {
    std::lock_guard<std::mutex> guard(self->numbers_mutex_);
    uint64_t new_total = self->total_bytes_allocated + size;
    if (new_total < size ||
        (self->total_cap_ && new_total > self->total_cap_)) {
      // Brittle "OOM" - even freeing memory won't undo it.
      self->total_bytes_allocated = self->total_cap_;
      self->seen_oom = true;
      return nullptr;
    }
    uint64_t new_bytes_in_use = self->bytes_in_use_ + size;
    if (new_bytes_in_use < size ||
        (self->cap_ && new_bytes_in_use > self->cap_)) {
      // Soft "OOM"
      self->seen_oom = true;
      return nullptr;
    }
    self->num_allocations_++;
    self->total_allocations++;
    self->bytes_in_use_ = new_bytes_in_use;
    self->max_bytes_in_use = std::max(self->max_bytes_in_use, new_bytes_in_use);
    self->total_bytes_allocated = new_total;
  }
  void* result = self->inner_->alloc(self->inner_->opaque, size);
  if (result != nullptr) {
    std::lock_guard<std::mutex> guard(self->map_mutex_);
    self->allocations_[result] = size;
  } else {
    // Hard "OOM" - rollback accounting.
    std::lock_guard<std::mutex> guard(self->numbers_mutex_);
    self->seen_oom = true;
    self->num_allocations_--;
    self->bytes_in_use_ -= size;
  }
  return result;
}

void TrackingMemoryManager::Free(void* opaque, void* address) {
  if (opaque == nullptr) {
    JXL_DEBUG_ABORT("Internal logic error");
    return;
  }
  if (address == nullptr) return;
  TrackingMemoryManager* self =
      reinterpret_cast<TrackingMemoryManager*>(opaque);
  bool found = false;
  size_t size = 0;
  {
    std::lock_guard<std::mutex> guard(self->map_mutex_);
    auto entry = self->allocations_.find(address);
    if (entry != self->allocations_.end()) {
      found = true;
      size = entry->second;
      self->allocations_.erase(entry);
    } else {
      JXL_DEBUG_ABORT("Internal logic error");
    }
  }

  if (found) {
    std::lock_guard<std::mutex> guard(self->numbers_mutex_);
    self->num_allocations_--;
    self->bytes_in_use_ -= size;
  }
  self->inner_->free(self->inner_->opaque, address);
}

jxl::Status TrackingMemoryManager::Reset() {
  if (num_allocations_ != 0) {
    return JXL_FAILURE("Memory leak");
  }
  if (!allocations_.empty()) {
    return JXL_FAILURE("Internal logic error");
  }
  if (bytes_in_use_ != 0) {
    return JXL_FAILURE("Internal logic error");
  }
  seen_oom = false;
  max_bytes_in_use = 0;
  total_allocations = 0;
  total_bytes_allocated = 0;
  return true;
}

}  // namespace tools
}  // namespace jpegxl