// Copyright 2024 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "user/user_cache.h"

#include <cstdlib>
#include <cstring>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <utility>

#include <mujoco/mjplugin.h>
#include "user/user_resource.h"


// makes a copy for user (strip unnecessary items)
mjCAsset mjCAsset::Copy(const mjCAsset& other) {
  mjCAsset asset;
  asset.id_ = other.Id();
  asset.timestamp_ = other.Timestamp();
  asset.data_ = other.data_;
  asset.size_ = other.size_;
  return asset;
}



// sets the total maximum size of the cache in bytes
// low-priority cached assets will be dropped to make the new memory
// requirement
void mjCCache::SetCapacity(std::size_t size) {
  std::lock_guard<std::mutex> lock(mutex_);
  capacity_ = size;
  Trim();
}



// returns the corresponding timestamp, if the given asset is stored in the cache
const std::string* mjCCache::HasAsset(const std::string& id) {
  std::lock_guard<std::mutex> lock(mutex_);
  auto it = lookup_.find(id);
  if (it == lookup_.end()) {
    return nullptr;
  }

  return &(it->second.Timestamp());
}



// inserts an asset into the cache, if asset is already in the cache, its data
// is updated only if the timestamps disagree
bool mjCCache::Insert(const std::string& modelname, const std::string& id, const mjResource *resource,
                      std::shared_ptr<const void> data, std::size_t size) {
  std::lock_guard<std::mutex> lock(mutex_);

  // check if asset is too large to fit in the cache
  if ((size_ + size > capacity_) &&
      lookup_.find(id) == lookup_.end()) {
    return false;
  }
  mjCAsset asset(modelname, id, resource, data, size);
  auto [it, inserted] = lookup_.insert({id, asset});
  mjCAsset* asset_ptr = &(it->second);

  if (!inserted) {
    if (size_ - asset_ptr->BytesCount() + size > capacity_) {
      return false;
    }
    models_[modelname].insert(asset_ptr);  // add it for the model
    asset_ptr->AddReference(modelname);
    if (it->second.Timestamp() == asset.Timestamp()) {
      return true;
    }
    asset_ptr->SetTimestamp(asset.Timestamp());
    size_ = size_ - asset_ptr->BytesCount() + size;
    asset_ptr->ReplaceData(asset);
    return true;
  }

  // new asset
  asset_ptr->SetInsertNum(insert_num_++);
  entries_.insert(asset_ptr);
  models_[modelname].insert(asset_ptr);
  size_ += size;
  return true;
}



// populate data from the cache into the given function, return true if data was
// copied
bool mjCCache::PopulateData(const std::string& id, const mjResource* resource, mjCDataFunc fn) {
  std::lock_guard<std::mutex> lock(mutex_);
  auto it = lookup_.find(id);
  if (it == lookup_.end()) {
    return false;
  }

  if (mju_isModifiedResource(resource, it->second.Timestamp().c_str())) {
    return false;
  }

  mjCAsset* asset = &(it->second);

  // update priority queue
  entries_.erase(asset);
  asset->IncrementAccess();
  entries_.insert(asset);

  return asset->PopulateData(fn);
}



// removes model from the cache along with assets referencing only this model
void mjCCache::RemoveModel(const std::string& filename) {
  std::lock_guard<std::mutex> lock(mutex_);
  for (mjCAsset* asset : models_[filename]) {
    asset->RemoveReference(filename);
    if (!asset->HasReferences()) {
      Delete(asset, filename);
    }
  }
  models_.erase(filename);
}



// Wipes out all internal data for the given model
void mjCCache::Reset(const std::string& filename) {
  std::lock_guard<std::mutex> lock(mutex_);
  for (auto asset : models_[filename]) {
    Delete(asset, filename);
  }
  models_.erase(filename);
}



// Wipes out all internal data
void mjCCache::Reset() {
  std::lock_guard<std::mutex> lock(mutex_);
  entries_.clear();
  lookup_.clear();
  models_.clear();
  size_ = 0;
  insert_num_ = 0;
}



std::size_t mjCCache::Capacity() const {
  std::lock_guard<std::mutex> lock(mutex_);
  return capacity_;
}



std::size_t mjCCache::Size() const {
  std::lock_guard<std::mutex> lock(mutex_);
  return size_;
}



// Deletes a single asset
void mjCCache::DeleteAsset(const std::string& id) {
  std::lock_guard<std::mutex> lock(mutex_);
  auto it = lookup_.find(id);
  if (it != lookup_.end()) {
    Delete(&(it->second));
  }
}



// Deletes a single asset (internal)
void mjCCache::Delete(mjCAsset* asset) {
  size_ -= asset->BytesCount();
  entries_.erase(asset);
  for (auto& reference : asset->References()) {
    models_[reference].erase(asset);
  }
  lookup_.erase(asset->Id());
}



// Deletes a single asset (internal)
void mjCCache::Delete(mjCAsset* asset, const std::string& skip) {
  size_ -= asset->BytesCount();
  entries_.erase(asset);

  for (auto& reference : asset->References()) {
    if (reference != skip) {
      models_[reference].erase(asset);
    }
  }
  lookup_.erase(asset->Id());
}



// trims out data to meet memory requirements
void mjCCache::Trim() {
  while (size_ > capacity_) {
    Delete(*entries_.begin());
  }
}