#include "onnx_helper.h"
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <gtest/gtest.h>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;

namespace {

bool is_aligned(const void *ptr, size_t align) {
  return align == 0 || (reinterpret_cast<uintptr_t>(ptr) % align) == 0;
}

int64_t get_external_i64(const TensorProto &tensor, const char *key) {
  for (const auto &entry : tensor.ref_external_data()) {
    if (entry.ref_key() == key) {
      return entry.ref_value().toint64();
    }
  }
  return -1;
}

std::string get_external_location(const TensorProto &tensor) {
  for (const auto &entry : tensor.ref_external_data()) {
    if (entry.ref_key() == "location") {
      return entry.ref_value();
    }
  }
  return "";
}

} // namespace

TEST(onnx_alignment_options, AlignmentMustBePowerOfTwo) {
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");
  TensorProto *t = graph->add_initializer();
  t->set_name("w0");
  t->set_data_type(TensorProto::DataType::FLOAT);
  t->ref_dims().push_back(2);
  t->ref_raw_data() = std::vector<uint8_t>{1, 2, 3, 4, 5, 6, 7, 8};

  SerializeOptions bad_sopts;
  bad_sopts.raw_data_threshold = 0;
  bad_sopts.alignment = 3;
  std::string serialized_model;
  std::unordered_map<std::string, std::string> external_files;
  EXPECT_THROW(
      model.SerializeToString(serialized_model, external_files, 1024, "weights", bad_sopts),
      std::runtime_error);

  TensorProto tensor;
  tensor.set_name("inline");
  tensor.set_data_type(TensorProto::DataType::FLOAT);
  tensor.ref_dims().push_back(2);
  tensor.ref_raw_data() = std::vector<uint8_t>{9, 10, 11, 12, 13, 14, 15, 16};
  std::string serialized_tensor;
  tensor.SerializeToString(serialized_tensor);

  ParseOptions bad_popts;
  bad_popts.alignment = 6;
  TensorProto parsed;
  EXPECT_THROW(parsed.ParseFromString(serialized_tensor, bad_popts), std::runtime_error);
}

TEST(onnx_alignment_options, SerializeToStringAlignmentSplitExternalFiles) {
  constexpr int64_t align = 16;
  constexpr int64_t max_external_file_size = 48;

  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  std::vector<std::vector<uint8_t>> payloads = {
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
      {12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
      {24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35},
      {36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47},
  };

  for (size_t i = 0; i < payloads.size(); ++i) {
    TensorProto *t = graph->add_initializer();
    t->set_name("w" + std::to_string(i));
    t->set_data_type(TensorProto::DataType::UINT8);
    t->ref_dims().push_back(static_cast<int64_t>(payloads[i].size()));
    t->ref_raw_data() = payloads[i];
  }

  SerializeOptions sopts;
  sopts.raw_data_threshold = 0;
  sopts.alignment = align;
  std::string serialized;
  std::unordered_map<std::string, std::string> external_files;
  model.SerializeToString(serialized, external_files, static_cast<size_t>(max_external_file_size),
                          "weights_part", sopts);

  EXPECT_GE(external_files.size(), 2u);

  ModelProto parsed;
  parsed.ParseFromString(serialized);
  ASSERT_EQ(parsed.ref_graph().ref_initializer().size(), payloads.size());

  for (size_t i = 0; i < payloads.size(); ++i) {
    const TensorProto &t = parsed.ref_graph().ref_initializer()[i];
    const std::string location = get_external_location(t);
    const int64_t off = get_external_i64(t, "offset");
    const int64_t len = get_external_i64(t, "length");
    ASSERT_FALSE(location.empty());
    ASSERT_GE(off, 0);
    ASSERT_EQ(off % align, 0) << "tensor " << i << " offset=" << off;
    ASSERT_EQ(len, static_cast<int64_t>(payloads[i].size()));
    ASSERT_EQ(external_files.count(location), 1u);
    const std::string &buffer = external_files.at(location);
    ASSERT_GE(static_cast<int64_t>(buffer.size()), off + len);
    EXPECT_EQ(std::memcmp(buffer.data() + off, payloads[i].data(), payloads[i].size()), 0);
  }
}

TEST(onnx_alignment_options, ParseAlignmentIncompatibleWithExternalDataWarnsOrErrorsNoCopy) {
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  const std::vector<float> vals0(4, 1.0f); // 16 bytes
  const std::vector<float> vals1(4, 2.0f); // 16 bytes
  auto add_tensor = [&](const std::string &name, const std::vector<float> &vals) {
    TensorProto *t = graph->add_initializer();
    t->set_name(name);
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->ref_dims().push_back(static_cast<int64_t>(vals.size()));
    t->ref_raw_data().resize(vals.size() * sizeof(float));
    std::memcpy(t->ref_raw_data().data(), vals.data(), vals.size() * sizeof(float));
  };
  add_tensor("w0", vals0);
  add_tensor("w1", vals1);

  const std::string onnx_file = "test_parse_alignment_incompatible.onnx";
  const std::string weights_file = "test_parse_alignment_incompatible.data";
  {
    utils::TwoFilesWriteStream wstream(onnx_file, weights_file);
    SerializeOptions sopts;
    sopts.raw_data_threshold = 0;
    sopts.alignment = 16;
    SerializeProtoToStream(model, wstream, sopts);
  }

  ModelProto loaded;
  std::string stderr_text;
  {
    testing::internal::CaptureStderr();
    utils::TwoFilesStream rstream(onnx_file, weights_file);
    ParseOptions ropts;
    ropts.alignment = 64;
    ParseProtoFromStream(loaded, rstream, ropts);
    stderr_text = testing::internal::GetCapturedStderr();
  }
  EXPECT_NE(stderr_text.find("Warning:"), std::string::npos);
  EXPECT_NE(stderr_text.find("incompatible with ParseOptions.alignment"), std::string::npos);
  ASSERT_EQ(loaded.ref_graph().ref_initializer().size(), 2u);
  EXPECT_TRUE(loaded.ref_graph().ref_initializer()[0].ref_raw_data().is_aligned_owned());
  EXPECT_TRUE(is_aligned(loaded.ref_graph().ref_initializer()[0].ref_raw_data().data(), 64));
  EXPECT_TRUE(loaded.ref_graph().ref_initializer()[1].ref_raw_data().is_aligned_owned());
  EXPECT_TRUE(is_aligned(loaded.ref_graph().ref_initializer()[1].ref_raw_data().data(), 64));

  ModelProto loaded_fail;
  {
    utils::TwoFilesStream rstream(onnx_file, weights_file);
    ParseOptions ropts;
    ropts.alignment = 64;
    ropts.no_copy = true;
    EXPECT_THROW(ParseProtoFromStream(loaded_fail, rstream, ropts), std::runtime_error);
  }

  std::remove(onnx_file.c_str());
  std::remove(weights_file.c_str());
}

TEST(onnx_alignment_options, ParseAlignmentNoCopyInlineRawDataRequiresAlignedOffset) {
  TensorProto tensor;
  tensor.set_name("inline");
  tensor.set_data_type(TensorProto::DataType::FLOAT);
  tensor.ref_dims().push_back(2);
  tensor.ref_raw_data() = std::vector<uint8_t>{1, 2, 3, 4, 5, 6, 7, 8};

  std::string serialized_tensor;
  ASSERT_NO_THROW(tensor.SerializeToString(serialized_tensor));

  ParseOptions ropts;
  ropts.alignment = 64;
  ropts.no_copy = true;

  TensorProto parsed;
  try {
    parsed.ParseFromString(serialized_tensor, ropts);
    FAIL() << "Expected ParseFromString to throw for incompatible no_copy alignment.";
  } catch (const std::runtime_error &e) {
    const std::string message = e.what();
    EXPECT_NE(message.find("incompatible with ParseOptions.alignment"), std::string::npos);
    EXPECT_NE(message.find("when no_copy=true"), std::string::npos);
  }
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferBasic) {
  // Build a model with three inline tensors of different sizes.
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  const std::vector<uint8_t> data0 = {1, 2, 3, 4, 5, 6, 7, 8};
  const std::vector<uint8_t> data1 = {9, 10, 11, 12};
  const std::vector<uint8_t> data2 = {13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24};

  for (size_t i = 0; i < 3; ++i) {
    TensorProto *t = graph->add_initializer();
    t->set_name("w" + std::to_string(i));
    t->set_data_type(TensorProto::DataType::UINT8);
  }
  graph->ref_initializer()[0].ref_raw_data() = data0;
  graph->ref_initializer()[1].ref_raw_data() = data1;
  graph->ref_initializer()[2].ref_raw_data() = data2;

  // Consolidate all tensors (threshold=0, override the default of 64).
  TensorBufferOptions opts;
  opts.raw_data_threshold = 0;
  auto buf = ConsolidateTensorsToBuffer(model, opts);
  ASSERT_NE(buf, nullptr);

  const ModelProto &cmodel = model;
  const auto &inits = cmodel.ref_graph().ref_initializer();

  // All tensors should now be borrowed.
  for (const TensorProto &t : inits) {
    EXPECT_TRUE(t.ref_raw_data().is_borrowed());
  }

  // Data must be correct.
  EXPECT_EQ(std::memcmp(inits[0].ref_raw_data().data(), data0.data(), data0.size()), 0);
  EXPECT_EQ(std::memcmp(inits[1].ref_raw_data().data(), data1.data(), data1.size()), 0);
  EXPECT_EQ(std::memcmp(inits[2].ref_raw_data().data(), data2.data(), data2.size()), 0);

  // All tensor data pointers should be within the same allocation.
  const uint8_t *base = buf.get();
  for (const TensorProto &t : inits) {
    EXPECT_GE(t.ref_raw_data().data(), base);
  }
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferWithAlignment) {
  constexpr size_t align = 16;
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  // Three tensors with payloads that are not multiples of 16.
  const std::vector<std::vector<uint8_t>> payloads = {
      std::vector<uint8_t>(7, 1),
      std::vector<uint8_t>(9, 2),
      std::vector<uint8_t>(13, 3),
  };
  for (size_t i = 0; i < payloads.size(); ++i) {
    TensorProto *t = graph->add_initializer();
    t->set_name("w" + std::to_string(i));
    t->set_data_type(TensorProto::DataType::UINT8);
    t->ref_dims().push_back(static_cast<int64_t>(payloads[i].size()));
    t->ref_raw_data() = payloads[i];
  }

  TensorBufferOptions opts;
  opts.alignment = static_cast<int64_t>(align);
  opts.raw_data_threshold = 0; // Override default 64: consolidate all tensors.
  auto buf = ConsolidateTensorsToBuffer(model, opts);
  ASSERT_NE(buf, nullptr);

  for (size_t i = 0; i < payloads.size(); ++i) {
    const TensorProto &t = graph->ref_initializer()[i];
    EXPECT_TRUE(t.ref_raw_data().is_borrowed());
    EXPECT_TRUE(is_aligned(t.ref_raw_data().data(), align))
        << "tensor " << i << " is not aligned to " << align;
    EXPECT_EQ(std::memcmp(t.ref_raw_data().data(), payloads[i].data(), payloads[i].size()), 0);
  }
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferThreshold) {
  // Only tensors >= threshold should be moved; smaller ones stay in owned mode.
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  const std::vector<uint8_t> small_data = {1, 2, 3}; // 3 bytes  < threshold=8
  const std::vector<uint8_t> large_data(16, 42);     // 16 bytes >= threshold=8

  TensorProto *t0 = graph->add_initializer();
  t0->set_name("small");
  t0->set_data_type(TensorProto::DataType::UINT8);
  t0->ref_raw_data() = small_data;

  TensorProto *t1 = graph->add_initializer();
  t1->set_name("large");
  t1->set_data_type(TensorProto::DataType::UINT8);
  t1->ref_raw_data() = large_data;

  TensorBufferOptions opts;
  opts.raw_data_threshold = 8;
  auto buf = ConsolidateTensorsToBuffer(model, opts);
  ASSERT_NE(buf, nullptr);

  const ModelProto &cmodel = model;
  const auto &inits = cmodel.ref_graph().ref_initializer();

  // Small tensor stays in owned (non-borrowed) mode.
  EXPECT_FALSE(inits[0].ref_raw_data().is_borrowed());
  // Large tensor is borrowed.
  EXPECT_TRUE(inits[1].ref_raw_data().is_borrowed());
  EXPECT_EQ(std::memcmp(inits[1].ref_raw_data().data(), large_data.data(), large_data.size()), 0);
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferLifetimeViaTensors) {
  // After the returned shared_ptr is reset, the tensor data must still be valid
  // because each tensor's ByteSpan holds its own shared_ptr owner token.
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  const std::vector<uint8_t> data = {10, 20, 30, 40};
  TensorProto *t = graph->add_initializer();
  t->set_name("w");
  t->set_data_type(TensorProto::DataType::UINT8);
  t->ref_raw_data() = data;

  TensorBufferOptions opts;
  opts.raw_data_threshold = 0; // Override default 64: consolidate small tensors too.
  std::shared_ptr<uint8_t[]> buf = ConsolidateTensorsToBuffer(model, opts);
  ASSERT_NE(buf, nullptr);
  buf.reset(); // Release caller's reference; buffer should stay alive via tensor.

  const ModelProto &cmodel = model;
  const TensorProto &ct = cmodel.ref_graph().ref_initializer()[0];
  EXPECT_TRUE(ct.ref_raw_data().is_borrowed());
  EXPECT_EQ(std::memcmp(ct.ref_raw_data().data(), data.data(), data.size()), 0);
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferEmptyModel) {
  ModelProto model;
  model.add_graph()->set_name("g");
  TensorBufferOptions opts;
  auto buf = ConsolidateTensorsToBuffer(model, opts);
  EXPECT_EQ(buf, nullptr);
}

TEST(onnx_alignment_options, ConsolidateTensorsToBufferInvalidAlignment) {
  ModelProto model;
  GraphProto *g = model.add_graph();
  g->set_name("g");
  TensorProto *t = g->add_initializer();
  t->set_data_type(TensorProto::DataType::UINT8);
  t->ref_raw_data() = std::vector<uint8_t>{1, 2, 3, 4};
  TensorBufferOptions opts;
  opts.alignment = 3; // Not a power of two – must throw.
  EXPECT_THROW(ConsolidateTensorsToBuffer(model, opts), std::runtime_error);
}

TEST(onnx_alignment_options, TensorBufferOptionsInheritedByParseOptions) {
  TensorBufferOptions base_opts;
  EXPECT_EQ(base_opts.raw_data_threshold, kSmallTensorDataThresholdBytes);
  EXPECT_EQ(base_opts.alignment, 0);

  ParseOptions popts;
  EXPECT_EQ(popts.raw_data_threshold, kSmallTensorDataThresholdBytes);
  EXPECT_EQ(popts.alignment, 0);
  // Verify the base-class fields are accessible through the derived type.
  TensorBufferOptions &base = popts;
  EXPECT_EQ(base.raw_data_threshold, kSmallTensorDataThresholdBytes);
}

TEST(onnx_alignment_options, TensorBufferOptionsInheritedBySerializeOptions) {
  SerializeOptions sopts;
  EXPECT_EQ(sopts.raw_data_threshold, kSmallTensorDataThresholdBytes);
  EXPECT_EQ(sopts.alignment, 0);
  TensorBufferOptions &base = sopts;
  EXPECT_EQ(base.raw_data_threshold, kSmallTensorDataThresholdBytes);
}

TEST(onnx_alignment_options, ParseNoCopyExternalDataRemainsValidAfterStreamDestruction) {
  constexpr int64_t align = 16;
  constexpr int64_t max_external_file_size = 32;
  const std::string onnx_file = "test_parse_no_copy_external_shared.onnx";
  const std::string weights_file = "test_parse_no_copy_external_shared.data";
  const std::string weights_file_1 = "test_parse_no_copy_external_shared.data.1";

  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");

  std::vector<std::vector<uint8_t>> payloads = {
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
      {12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23},
      {24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35},
      {36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47},
  };

  for (size_t i = 0; i < payloads.size(); ++i) {
    TensorProto *t = graph->add_initializer();
    t->set_name("w" + std::to_string(i));
    t->set_data_type(TensorProto::DataType::UINT8);
    t->ref_dims().push_back(static_cast<int64_t>(payloads[i].size()));
    t->ref_raw_data() = payloads[i];
  }

  {
    utils::TwoFilesWriteStream wstream(onnx_file, weights_file);
    SerializeOptions sopts;
    sopts.raw_data_threshold = 0;
    sopts.alignment = align;
    sopts.max_external_file_size = max_external_file_size;
    SerializeProtoToStream(model, wstream, sopts);
  }

  ModelProto metadata;
  {
    utils::FileStream meta_stream(onnx_file);
    ParseOptions meta_opts;
    ParseProtoFromStream(metadata, meta_stream, meta_opts, false);
  }

  ModelProto loaded;
  {
    utils::TwoFilesStream rstream(onnx_file, weights_file);
    ParseOptions ropts;
    ropts.no_copy = true;
    ropts.alignment = align;
    ParseProtoFromStream(loaded, rstream, ropts);
  }
  // rstream is destroyed here; borrowed raw_data must remain valid via ByteSpan shared ownership.

  ASSERT_EQ(loaded.ref_graph().ref_initializer().size(), payloads.size());
  std::unordered_map<std::string, const uint8_t *> base_ptrs;
  for (size_t i = 0; i < payloads.size(); ++i) {
    const TensorProto &meta = metadata.ref_graph().ref_initializer()[i];
    const TensorProto &t = loaded.ref_graph().ref_initializer()[i];
    ASSERT_TRUE(t.ref_raw_data().is_borrowed());
    ASSERT_TRUE(is_aligned(t.ref_raw_data().data(), align));
    EXPECT_EQ(std::memcmp(t.ref_raw_data().data(), payloads[i].data(), payloads[i].size()), 0);
    const std::string location = get_external_location(meta);
    const int64_t offset = get_external_i64(meta, "offset");
    ASSERT_FALSE(location.empty());
    ASSERT_GE(offset, 0);
    ASSERT_EQ(offset % align, 0);
    const uint8_t *base = t.ref_raw_data().data() - offset;
    auto inserted = base_ptrs.emplace(location, base);
    if (!inserted.second) {
      EXPECT_EQ(inserted.first->second, base);
    }
  }

  EXPECT_GE(base_ptrs.size(), 2u);
  std::unordered_set<const uint8_t *> distinct_bases;
  for (const auto &item : base_ptrs) {
    distinct_bases.insert(item.second);
  }
  EXPECT_EQ(distinct_bases.size(), base_ptrs.size());

  std::remove(onnx_file.c_str());
  std::remove(weights_file.c_str());
  std::remove(weights_file_1.c_str());
}

namespace {

// Builds a model with a single initializer already marked EXTERNAL and carrying the given offset.
// The inline raw_data is preserved so it can be written out during serialization;
// PopulateExternalData skips tensors that are already external, so the crafted offset reaches the
// write path unchanged.
ModelProto MakeModelWithExternalOffset(const std::string &location, int64_t offset) {
  ModelProto model;
  GraphProto *graph = model.add_graph();
  graph->set_name("g");
  TensorProto *t = graph->add_initializer();
  t->set_name("w0");
  t->set_data_type(TensorProto::DataType::UINT8);
  t->ref_dims().push_back(4);
  t->ref_raw_data() = std::vector<uint8_t>{1, 2, 3, 4};
  t->ref_data_location() = TensorProto::DataLocation::EXTERNAL;
  StringStringEntryProto *loc = t->add_external_data();
  loc->set_key("location");
  loc->set_value(location);
  StringStringEntryProto *off = t->add_external_data();
  off->set_key("offset");
  off->set_value(std::to_string(offset));
  StringStringEntryProto *len = t->add_external_data();
  len->set_key("length");
  len->set_value(std::to_string(t->ref_raw_data().size()));
  return model;
}

} // namespace

// Propagates onnx/onnx#8260: a crafted external-data offset that would pad the weights file with an
// unbounded amount of zeros is rejected during serialization instead of silently writing the zeros.
TEST(onnx_alignment_options, SerializeRejectsUnboundedExternalDataPadding) {
  const std::string onnx_file = "test_external_padding_reject.onnx";
  const std::string weights_file = "test_external_padding_reject.data";
  // The tensor sits at offset 0, so the padding equals the offset. One byte over the 64 KiB cap is
  // the exact boundary complement of SerializeAcceptsBoundedExternalDataPadding (which uses 64
  // KiB).
  ModelProto model = MakeModelWithExternalOffset(weights_file, 64 * 1024 + 1);

  utils::TwoFilesWriteStream wstream(onnx_file, weights_file);
  SerializeOptions sopts;
  sopts.raw_data_threshold = 0;
  sopts.alignment = 16;
  EXPECT_THROW(
      {
        try {
          SerializeProtoToStream(model, wstream, sopts);
        } catch (const std::runtime_error &err) {
          EXPECT_NE(std::string(err.what()).find("exceeds the maximum"), std::string::npos)
              << err.what();
          throw;
        }
      },
      std::runtime_error);

  std::remove(onnx_file.c_str());
  std::remove(weights_file.c_str());
}

// A padding gap up to the cap is legitimate (e.g. an aligned offset) and must still serialize.
TEST(onnx_alignment_options, SerializeAcceptsBoundedExternalDataPadding) {
  const std::string onnx_file = "test_external_padding_accept.onnx";
  const std::string weights_file = "test_external_padding_accept.data";
  // 64 KiB offset equals the cap, so the padding stays within bounds.
  ModelProto model = MakeModelWithExternalOffset(weights_file, 64 * 1024);

  {
    utils::TwoFilesWriteStream wstream(onnx_file, weights_file);
    SerializeOptions sopts;
    sopts.raw_data_threshold = 0;
    sopts.alignment = 16;
    EXPECT_NO_THROW(SerializeProtoToStream(model, wstream, sopts));
  }

  std::remove(onnx_file.c_str());
  std::remove(weights_file.c_str());
}
