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

// Tests for user/user_resource.cc

#include "src/user/user_resource.h"

#include <array>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <string>
#include <utility>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <absl/strings/str_format.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_plugin.h"
#include "src/engine/engine_util_misc.h"
#include "test/fixture.h"

namespace mujoco {
namespace {

using ::testing::HasSubstr;
using ::testing::IsNull;
using ::testing::NotNull;
using ::testing::StrEq;

using ResourceTest = MujocoTest;

int open_nop(mjResource* resource) {
  return 1;
}

int open_str(mjResource* resource) {
  if (std::strcmp(resource->name, "str:file")) {
    return 0;
  }

  const std::size_t kBufferSize = 100;
  resource->data = mju_malloc(kBufferSize * sizeof(char));
  absl::SNPrintF((char*) resource->data, kBufferSize, "Hello World");
  return 1;
}

int read_nop(mjResource* resource, const void** buffer) {
  return 0;
}

int read_str(mjResource* resource, const void** buffer) {
  *buffer = resource->data;
  return std::strlen((const char*) resource->data) + 1;
}

void close_nop(mjResource* resource) {
}

void close_str(mjResource* resource) {
  mju_free(resource->data);
  resource->data = nullptr;
}

TEST_F(ResourceTest, RegisterProviderSuccess) {
  mjpResourceProvider provider = {
      .prefix = "my-prefix.123+45",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  int count1 = mjp_resourceProviderCount();
  int i = mjp_registerResourceProvider(&provider);
  int count2 = mjp_resourceProviderCount();

  EXPECT_GT(i, 0);
  EXPECT_EQ(count2 - count1, 1);
}

TEST_F(ResourceTest, RegisterProviderMultipleSuccess) {
  mjpResourceProvider provider = {
      .prefix = "my-prefix.123+44",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  mjpResourceProvider provider2 = {
      .prefix = "my-prefix.123+46",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  mjpResourceProvider provider3 = {
      .prefix = "my-prefix.123+41",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  int count1 = mjp_resourceProviderCount();
  int i = mjp_registerResourceProvider(&provider);
  int i2 = mjp_registerResourceProvider(&provider2);
  int i3 = mjp_registerResourceProvider(&provider3);
  int count2 = mjp_resourceProviderCount();

  EXPECT_GT(i, 0);
  EXPECT_GT(i2, 0);
  EXPECT_GT(i3, 0);
  EXPECT_EQ(count2 - count1, 3);
}

TEST_F(ResourceTest, RegisterProviderMissingCallbacks) {
  mjpResourceProvider provider = {
      .prefix = "myprefix",
  };

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings("callback");

  int i = mjp_registerResourceProvider(&provider);
  EXPECT_LT(i, 1);
}

TEST_F(ResourceTest, RegisterProviderMissingPrefix) {
  mjpResourceProvider provider = {
      .prefix = "",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings("prefix");

  int i = mjp_registerResourceProvider(&provider);
  EXPECT_LT(i, 1);
}

TEST_F(ResourceTest, RegisterProviderInvalidPrefix1) {
  mjpResourceProvider provider = {
      .prefix = "1invalid",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings("prefix");

  int i = mjp_registerResourceProvider(&provider);
  EXPECT_LT(i, 1);
}

TEST_F(ResourceTest, RegisterProviderInvalidPrefix2) {
  mjpResourceProvider provider = {
      .prefix = "invalid:",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings("prefix");

  int i = mjp_registerResourceProvider(&provider);
  EXPECT_LT(i, 1);
}

TEST_F(ResourceTest, RegisterProviderSame) {
  mjpResourceProvider provider = {
      .prefix = "prefix",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  mjpResourceProvider provider2 = {
      .prefix = "prefix",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  int i1 = mjp_registerResourceProvider(&provider);
  int count1 = mjp_resourceProviderCount();

  int i2 = mjp_registerResourceProvider(&provider2);
  int count2 = mjp_resourceProviderCount();

  EXPECT_EQ(i1, i2);
  EXPECT_EQ(count1, count2);
}

TEST_F(ResourceTest, RegisterProviderSameCase) {
  mjpResourceProvider provider = {
      .prefix = "prefix",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  mjpResourceProvider provider2 = {
      .prefix = "PREFIX",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  int i1 = mjp_registerResourceProvider(&provider);
  int count1 = mjp_resourceProviderCount();

  int i2 = mjp_registerResourceProvider(&provider2);
  int count2 = mjp_resourceProviderCount();

  EXPECT_EQ(i1, i2);
  EXPECT_EQ(count1, count2);
}

TEST_F(ResourceTest, GeneralTest) {
  mjpResourceProvider provider = {
      .prefix = "str",
      .open = open_str,
      .read = read_str,
      .close = close_str,
  };

  // register resource provider
  int i = mjp_registerResourceProvider(&provider);
  EXPECT_GT(i, 0);

  // open resource
  mjResource* resource = mju_openResource("", "str:file", nullptr, nullptr, 0);
  ASSERT_THAT(resource,  NotNull());

  const char* buffer = NULL;
  int bytes = mju_readResource(resource, (const void**) &buffer);
  EXPECT_EQ(bytes, std::strlen("Hello World") + 1);
  EXPECT_THAT(buffer, StrEq("Hello World"));

  mju_closeResource(resource);
}

TEST_F(ResourceTest, GeneralFailureTest) {
  mjpResourceProvider provider = {
      .prefix = "str",
      .open = open_str,
      .read = read_str,
      .close = close_str,
  };

  // register resource provider
  int i = mjp_registerResourceProvider(&provider);
  EXPECT_GT(i, 0);

  static std::array<char, 1024> error;

  // open resource
  mjResource* resource = mju_openResource("", "str:notfound", nullptr,
                                          error.data(), error.size());
  ASSERT_THAT(resource, IsNull());

  EXPECT_THAT(error.data(), HasSubstr("Error opening file"));
}

TEST_F(ResourceTest, NameWithValidPrefix) {
  mjpResourceProvider provider = {
      .prefix = "nop",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  // register resource provider
  int i = mjp_registerResourceProvider(&provider);
  EXPECT_GT(i, 0);

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings();

  // open resource
  mjResource* resource = mju_openResource("", "nop:found", nullptr, nullptr, 0);
  ASSERT_THAT(resource, NotNull());
  mju_closeResource(resource);
}

TEST_F(ResourceTest, NameWithUpperCasePrefix) {
  mjpResourceProvider provider = {
      .prefix = "nop",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  // register resource provider
  int i = mjp_registerResourceProvider(&provider);
  EXPECT_GT(i, 0);

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings();

  // open resource
  mjResource* resource = mju_openResource("", "NOP:found", nullptr, nullptr, 0);
  ASSERT_THAT(resource, NotNull());
  mju_closeResource(resource);
}

TEST_F(ResourceTest, NameWithInvalidPrefix) {
  mjpResourceProvider provider = {
      .prefix = "nop",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  // register resource provider
  int i = mjp_registerResourceProvider(&provider);
  EXPECT_GT(i, 0);

  MockWarningHandler warning_handler;
  warning_handler.ExpectWarnings();

  // open resource
  mjResource* resource = mju_openResource("", "nopfound", nullptr, nullptr, 0);
  ASSERT_THAT(resource, IsNull());
}

TEST_F(ResourceTest, GetResourceDir) {
  const std::vector<std::pair<std::string, std::string>> cases = {
    { "foo/bar/baz", "foo/bar/" },
    { "foo/bar/", "foo/bar/" },
    { "foo/bar", "foo/" },
    { "/foo/bar", "/foo/" },
    { "/foo", "/" },
    { "/", "/" },
    { "", "" },
  };

  const char* dir = nullptr;
  int ndir = 0;

  mjpResourceProvider provider;
  provider.prefix = "provider";

  for (const auto& [name, expected] : cases) {
    mjResource resource;
    resource.provider = nullptr;
    resource.name = const_cast<char*>(name.c_str());
    mju_getResourceDir(&resource, &dir, &ndir);
    EXPECT_THAT(std::string(dir, ndir), expected);
  }
}

TEST_F(ResourceTest, GetResourceDirProvider) {
  const std::vector<std::pair<std::string, std::string>> cases = {
    { "provider:/foo/bar", "provider:/foo/" },
    { "provider:/foo/", "provider:/foo/" },
    { "provider:foo/", "provider:foo/" },
    { "provider:foo", "provider:" },
    { "provider:/", "provider:/" },
  };

  const char* dir = nullptr;
  int ndir = 0;

  mjpResourceProvider provider;
  provider.prefix = "provider";

  for (const auto& [name, expected] : cases) {
    mjResource resource;
    resource.provider = &provider;
    resource.name = const_cast<char*>(name.c_str());
    mju_getResourceDir(&resource, &dir, &ndir);
    EXPECT_THAT(std::string(dir, ndir), expected);
  }
}

TEST_F(ResourceTest, GetResourceDirNullResource) {
  const char* dir = nullptr;
  int ndir = 0;

  mju_getResourceDir(nullptr, &dir, &ndir);
  EXPECT_THAT(dir, IsNull());
  EXPECT_EQ(ndir, 0);
}

TEST_F(ResourceTest, GetResourceDirNullName) {
  const char* dir = nullptr;
  int ndir = 0;

  mjResource resource;
  resource.name = nullptr;
  resource.provider = nullptr;

  mju_getResourceDir(nullptr, &dir, &ndir);
  EXPECT_THAT(dir, IsNull());
  EXPECT_EQ(ndir, 0);
}

TEST_F(ResourceTest, OSFilesystemTimestamps) {
  time_t t;

  // some random file
  const char* const file = "engine/testdata/collision_box/boxbox_deep.xml";
  const std::string xml_path = GetTestDataFilePath(file);

  mjResource* resource = mju_openResource("", xml_path.c_str(), nullptr,
                                          nullptr, 0);
  mju_decodeBase64((uint8_t*) &t, resource->timestamp);

  // equal timestamps
  EXPECT_EQ(mju_isModifiedResource(resource, resource->timestamp), 0);

  std::array<char, 512> test_timestamp;

  // older resource timestamp
  t++;
  mju_encodeBase64(test_timestamp.data(), (uint8_t*) &t, sizeof(time_t));
  EXPECT_EQ(mju_isModifiedResource(resource, test_timestamp.data()), -1);


  // newer resource timestamp
  t -= 2;
  mju_encodeBase64(test_timestamp.data(), (uint8_t*) &t, sizeof(time_t));
  EXPECT_EQ(mju_isModifiedResource(resource, test_timestamp.data()), 1);

  mju_closeResource(resource);
}

// ===================== Write Resource Tests =====================

struct WriteBuffer {
  std::vector<uint8_t> data;
};

// static storage for the last written buffer (for round-trip testing)
static std::vector<uint8_t> last_written;

mjtSize write_capture(mjResource* resource, const void* buffer, mjtSize nbytes) {
  last_written.clear();
  const uint8_t* bytes = static_cast<const uint8_t*>(buffer);
  last_written.insert(last_written.end(), bytes, bytes + nbytes);
  return nbytes;
}

// open callback for reading back captured data
int open_read_capture(mjResource* resource) { return 1; }

int read_capture(mjResource* resource, const void** buffer) {
  *buffer = last_written.data();
  return static_cast<int>(last_written.size());
}

void close_read_capture(mjResource* resource) {}

TEST_F(ResourceTest, WriteResourceWithProvider) {
  // Register a provider with write callbacks
  mjpResourceProvider provider = {
      .prefix = "wrtest",
      .open = open_read_capture,
      .read = read_capture,
      .close = close_read_capture,
      .write = write_capture,
  };

  mjp_registerResourceProvider(&provider);

  // Write some data
  const char* data = "Hello, Resource Writer!";
  mjtSize nbytes = static_cast<mjtSize>(std::strlen(data));

  mjtSize written = mju_writeResource("wrtest:myfile.txt", data, nbytes, nullptr, nullptr, 0);
  EXPECT_EQ(written, nbytes);

  // Read it back via the same provider
  char error[256] = {0};
  mjResource* resource =
      mju_openResource("", "wrtest:myfile.txt", nullptr, error, sizeof(error));
  ASSERT_THAT(resource, NotNull());

  const void* buf = nullptr;
  int read_bytes = mju_readResource(resource, &buf);
  EXPECT_EQ(read_bytes, nbytes);
  EXPECT_EQ(std::memcmp(buf, data, nbytes), 0);

  mju_closeResource(resource);
}

TEST_F(ResourceTest, WriteResourceDefaultPosix) {
  // Write to a temp file using the default POSIX provider (no prefix)
  std::string tmpfile = testing::TempDir() + "/mj_write_test.bin";
  const char* data = "MuJoCo resource write test";
  mjtSize nbytes = static_cast<mjtSize>(std::strlen(data));

  mjtSize written = mju_writeResource(tmpfile.c_str(), data, nbytes, nullptr, nullptr, 0);
  EXPECT_EQ(written, nbytes);

  // Read it back via mju_openResource (default POSIX provider)
  char error[256] = {0};
  mjResource* resource =
      mju_openResource("", tmpfile.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(resource, NotNull());

  const void* buf = nullptr;
  int read_bytes = mju_readResource(resource, &buf);
  EXPECT_EQ(read_bytes, nbytes);
  EXPECT_EQ(std::memcmp(buf, data, nbytes), 0);

  mju_closeResource(resource);

  // Clean up
  std::remove(tmpfile.c_str());
}

TEST_F(ResourceTest, WriteResourceNoWriteCallback) {
  // Register a read-only provider (no write callback)
  mjpResourceProvider provider = {
      .prefix = "rdonly",
      .open = open_nop,
      .read = read_nop,
      .close = close_nop,
  };

  mjp_registerResourceProvider(&provider);

  const char* data = "data";
  mjtSize written = mju_writeResource("rdonly:somefile", data, 4, nullptr, nullptr, 0);
  EXPECT_EQ(written, -1);
}

}  // namespace
}  // namespace mujoco