// C++ conformance test for the native Shard protocol.
//
// This test does not check that C++ can round-trip its own output — that would
// only prove C++ is self-consistent. It parses the *Python-produced* committed
// vectors, asserts every field the protocol promises to carry, independently
// recomputes the CRC32C over the reassembled tensor, and re-serializes the
// message back out for the Python test to compare byte-for-byte. Only that
// closes the loop: both languages agree on the same bytes and the same meaning.
//
// Run via ctest; see packages/node/native/CMakeLists.txt.

#include "shard_runtime.pb.h"

#include <cstdint>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

// `pb` is already taken by protobuf's own generated headers.
namespace sp = meshnet::shard::v1;

namespace {

int g_failures = 0;

#define CHECK(cond)                                                            \
  do {                                                                         \
    if (!(cond)) {                                                             \
      std::cerr << "FAIL " << __FILE__ << ":" << __LINE__ << ": " << #cond     \
                << "\n";                                                       \
      ++g_failures;                                                            \
    }                                                                          \
  } while (0)

#define CHECK_EQ(actual, expected)                                             \
  do {                                                                         \
    auto &&a_ = (actual);                                                      \
    auto &&e_ = (expected);                                                    \
    if (!(a_ == e_)) {                                                         \
      std::cerr << "FAIL " << __FILE__ << ":" << __LINE__ << ": " << #actual   \
                << " == " << #expected << "\n  actual:   " << a_               \
                << "\n  expected: " << e_ << "\n";                             \
      ++g_failures;                                                            \
    }                                                                          \
  } while (0)

// Independent CRC32C (Castagnoli). Written from the polynomial rather than
// shared with the Python side on purpose: a checksum that both languages
// compute with the *same* code proves nothing about interoperability.
uint32_t Crc32c(const std::string &data) {
  static uint32_t table[256];
  static bool built = false;
  if (!built) {
    for (uint32_t i = 0; i < 256; ++i) {
      uint32_t c = i;
      for (int k = 0; k < 8; ++k) {
        c = (c & 1) ? (c >> 1) ^ 0x82F63B78u : (c >> 1);
      }
      table[i] = c;
    }
    built = true;
  }
  uint32_t crc = 0xFFFFFFFFu;
  for (unsigned char byte : data) {
    crc = (crc >> 8) ^ table[(crc ^ byte) & 0xFF];
  }
  return crc ^ 0xFFFFFFFFu;
}

std::string ReadFile(const std::filesystem::path &path) {
  std::ifstream in(path, std::ios::binary);
  if (!in) {
    std::cerr << "FAIL cannot read " << path << "\n";
    ++g_failures;
    return {};
  }
  std::ostringstream buffer;
  buffer << in.rdbuf();
  return buffer.str();
}

// Reassemble a tensor's fragments and validate coverage, exactly as a worker
// must before it feeds the bytes to a forward pass.
std::string ReassembleUncompressed(const sp::NamedTensor &tensor) {
  std::string body;
  std::vector<const sp::TensorFragment *> fragments;
  for (const auto &fragment : tensor.fragments()) {
    fragments.push_back(&fragment);
  }
  CHECK(!fragments.empty());
  for (uint32_t index = 0; index < fragments.size(); ++index) {
    const sp::TensorFragment *found = nullptr;
    for (const auto *fragment : fragments) {
      if (fragment->fragment_index() == index) {
        found = fragment;
        break;
      }
    }
    CHECK(found != nullptr);
    if (found == nullptr) {
      return {};
    }
    CHECK_EQ(found->fragment_count(),
             static_cast<uint32_t>(fragments.size()));
    // Offsets must tile the wire body exactly: no hole, no overlap.
    CHECK_EQ(found->byte_offset(), static_cast<uint64_t>(body.size()));
    body.append(found->payload());
  }
  return body;
}

void CheckFingerprint(const sp::Fingerprint &fingerprint) {
  CHECK_EQ(fingerprint.model_artifact_digest(), std::string("sha256:1f0c9d2e"));
  CHECK_EQ(fingerprint.runtime_recipe_digest(), std::string("sha256:ab77e410"));
  CHECK_EQ(fingerprint.recipe_id(), std::string("llama-gguf-q4km-rocm"));
  CHECK_EQ(fingerprint.recipe_version(), std::string("3"));
  CHECK_EQ(fingerprint.catalogue_version(), std::string("2026.07.1"));
}

// The canonical session request, as produced by Python.
void TestSessionRequestVector(const std::string &bytes) {
  sp::SessionRequest request;
  CHECK(request.ParseFromString(bytes));
  CHECK(request.kind_case() == sp::SessionRequest::kChunk);

  const sp::ActivationChunk &chunk = request.chunk();
  const sp::Envelope &envelope = chunk.envelope();

  // Every field the protocol promises to carry (acceptance criterion 4).
  CHECK_EQ(envelope.schema_version(), sp::SCHEMA_VERSION_1);
  CHECK_EQ(envelope.work_id(), std::string("work-7f3a"));
  CHECK_EQ(envelope.route_session_id(), std::string("rs-2b91"));
  CHECK_EQ(envelope.route_epoch(), 7u);
  CheckFingerprint(envelope.fingerprint());
  CHECK_EQ(envelope.shard_range().start_layer(), 12u);
  CHECK_EQ(envelope.shard_range().end_layer(), 24u);
  // Overlap-safe effective start (ADR-0012): a worker that begins at
  // start_layer instead would re-apply layers 12..15.
  CHECK_EQ(envelope.shard_range().effective_start_layer(), 16u);
  CHECK_EQ(envelope.phase(), sp::PHASE_PREFILL);
  CHECK_EQ(envelope.position().first_position(), 256u);
  CHECK_EQ(envelope.position().token_count(), 128u);
  CHECK_EQ(envelope.idempotency_step(), 42u);
  CHECK_EQ(envelope.cache_expectation().mode(), sp::CACHE_MODE_PREFILL);
  CHECK_EQ(envelope.cache_expectation().expected_past_len(), 256u);
  CHECK_EQ(envelope.deadline_unix_nanos(), 1800000000000000000LL);
  CHECK_EQ(envelope.chunk().chunk_index(), 1u);
  CHECK_EQ(envelope.chunk().chunk_count(), 3u);
  CHECK_EQ(envelope.chunk().final_chunk(), false);

  // The versioned named-tensor bundle (acceptance criterion 5).
  const sp::TensorBundle &bundle = chunk.bundle();
  CHECK_EQ(bundle.bundle_version(), 1u);
  CHECK_EQ(bundle.tensors_size(), 1);

  const sp::NamedTensor &tensor = bundle.tensors(0);
  CHECK_EQ(tensor.name(), std::string("hidden_states"));
  CHECK_EQ(tensor.shape_size(), 3);
  CHECK_EQ(tensor.shape(0), 1);
  CHECK_EQ(tensor.shape(1), 128);
  CHECK_EQ(tensor.shape(2), 8);
  CHECK_EQ(tensor.dtype(), sp::DTYPE_BFLOAT16);
  CHECK_EQ(tensor.byte_order(), sp::BYTE_ORDER_LITTLE_ENDIAN);
  CHECK_EQ(tensor.total_bytes(), 1u * 128u * 8u * 2u);
  CHECK_EQ(tensor.compression(), sp::COMPRESSION_NONE);

  // Multi-fragment on purpose: reassembly is what a worker actually does.
  CHECK(tensor.fragments_size() > 1);
  const std::string payload = ReassembleUncompressed(tensor);
  CHECK_EQ(payload.size(), static_cast<size_t>(tensor.total_bytes()));

  // The payload Python generated, recomputed here from its rule.
  std::string expected;
  expected.reserve(payload.size());
  for (size_t i = 0; i < payload.size(); ++i) {
    expected.push_back(static_cast<char>((i * 7 + 11) % 256));
  }
  CHECK(payload == expected);

  // Checksum: computed by Python, verified by an independent C++ CRC32C.
  CHECK_EQ(tensor.checksum().algorithm(), sp::CHECKSUM_ALGORITHM_CRC32C);
  const std::string &checksum = tensor.checksum().value();
  CHECK_EQ(checksum.size(), 4u);
  if (checksum.size() == 4) {
    const uint32_t actual = Crc32c(payload);
    const uint32_t declared =
        (static_cast<uint32_t>(static_cast<unsigned char>(checksum[0])) << 24) |
        (static_cast<uint32_t>(static_cast<unsigned char>(checksum[1])) << 16) |
        (static_cast<uint32_t>(static_cast<unsigned char>(checksum[2])) << 8) |
        static_cast<uint32_t>(static_cast<unsigned char>(checksum[3]));
    CHECK_EQ(actual, declared);
  }
}

void TestCapabilityReportVector(const std::string &bytes) {
  sp::CapabilityReport report;
  CHECK(report.ParseFromString(bytes));
  CHECK_EQ(report.schema_version(), sp::SCHEMA_VERSION_1);
  CheckFingerprint(report.fingerprint());
  CHECK_EQ(report.backend(), std::string("rocm"));
  CHECK_EQ(report.device(), std::string("gfx1151"));
  CHECK_EQ(report.validated(), true);
  CHECK_EQ(report.max_concurrent_sessions(), 4u);
  CHECK_EQ(report.max_context_tokens(), 8192u);
  CHECK_EQ(report.flow_control().max_prefill_chunk_tokens(), 128u);
  CHECK_EQ(report.flow_control().max_chunk_bytes(), 4u * 1024u * 1024u);
  CHECK_EQ(report.accepted_compression_size(), 2);
  CHECK_EQ(report.supported_schema_versions_size(), 1);
}

// Forward compatibility: a field this build has never heard of must survive a
// parse/serialize cycle untouched. Without this, an old Shard silently strips
// fields a newer peer depends on as it forwards activations down the route.
void TestUnknownFieldsArePreserved(const std::string &bytes) {
  // Field 9999, wire type 0 (varint), value 12345 — not in this schema.
  std::string forward = bytes;
  forward.push_back(static_cast<char>(0xB8));  // tag: (9999 << 3) | 0
  forward.push_back(static_cast<char>(0xE0));
  forward.push_back(static_cast<char>(0x04));
  forward.push_back(static_cast<char>(0xB9));  // value: 12345
  forward.push_back(static_cast<char>(0x60));

  sp::SessionRequest request;
  CHECK(request.ParseFromString(forward));
  // Known fields still parse.
  CHECK_EQ(request.chunk().envelope().work_id(), std::string("work-7f3a"));
  // And the unknown field is retained rather than dropped.
  CHECK(!request.GetReflection()->GetUnknownFields(request).empty());

  std::string reserialized;
  CHECK(request.SerializeToString(&reserialized));
  CHECK_EQ(reserialized.size(), forward.size());

  sp::SessionRequest reparsed;
  CHECK(reparsed.ParseFromString(reserialized));
  CHECK(!reparsed.GetReflection()->GetUnknownFields(reparsed).empty());
}

// Backward compatibility: a message from a peer that sets none of the optional
// groups must still parse, yielding proto3 defaults rather than an error.
void TestSparseMessageParses() {
  sp::SessionRequest minimal;
  minimal.mutable_chunk()->mutable_envelope()->set_work_id("w");
  std::string bytes;
  CHECK(minimal.SerializeToString(&bytes));

  sp::SessionRequest parsed;
  CHECK(parsed.ParseFromString(bytes));
  CHECK_EQ(parsed.chunk().envelope().work_id(), std::string("w"));
  CHECK_EQ(parsed.chunk().envelope().route_epoch(), 0u);
  CHECK_EQ(parsed.chunk().envelope().phase(), sp::PHASE_UNSPECIFIED);
  CHECK_EQ(parsed.chunk().bundle().tensors_size(), 0);
}

// The decode fast path must stay small: repeating the full envelope per token
// is pure overhead on the hottest path in the system.
void TestDecodeFastPathIsSmall() {
  sp::SessionRequest decode;
  sp::DecodeStep *step = decode.mutable_decode();
  step->set_idempotency_step(9);
  step->set_position(1024);
  step->set_expected_past_len(1024);
  step->set_work_id("work-7f3a");
  sp::NamedTensor *tensor = step->mutable_tensor();
  tensor->set_name("hidden_states");
  tensor->add_shape(1);
  tensor->add_shape(1);
  tensor->add_shape(8);
  tensor->set_dtype(sp::DTYPE_BFLOAT16);
  tensor->set_byte_order(sp::BYTE_ORDER_LITTLE_ENDIAN);
  tensor->set_total_bytes(16);
  tensor->set_compression(sp::COMPRESSION_NONE);
  sp::TensorFragment *fragment = tensor->add_fragments();
  fragment->set_fragment_index(0);
  fragment->set_fragment_count(1);
  fragment->set_byte_offset(0);
  fragment->set_payload(std::string(16, '\x01'));

  std::string bytes;
  CHECK(decode.SerializeToString(&bytes));
  // Payload is 16 bytes; the framing around it must stay well under the
  // envelope-carrying prefill path.
  CHECK(bytes.size() < 96);

  sp::SessionRequest parsed;
  CHECK(parsed.ParseFromString(bytes));
  CHECK(parsed.kind_case() == sp::SessionRequest::kDecode);
  CHECK_EQ(parsed.decode().position(), 1024u);
  CHECK_EQ(parsed.decode().tensor().total_bytes(), 16u);
}

}  // namespace

int main(int argc, char **argv) {
  GOOGLE_PROTOBUF_VERIFY_VERSION;
  if (argc < 3) {
    std::cerr << "usage: " << argv[0] << " <testdata-dir> <output-dir>\n";
    return 2;
  }
  const std::filesystem::path testdata(argv[1]);
  const std::filesystem::path output(argv[2]);

  const std::string session_bytes =
      ReadFile(testdata / "session_request_golden.binpb");
  const std::string capability_bytes =
      ReadFile(testdata / "capability_report_golden.binpb");

  TestSessionRequestVector(session_bytes);
  TestCapabilityReportVector(capability_bytes);
  TestUnknownFieldsArePreserved(session_bytes);
  TestSparseMessageParses();
  TestDecodeFastPathIsSmall();

  // Re-serialize the canonical message from the C++ object model and hand it
  // back for Python to compare against the golden bytes. If the two languages
  // disagreed about any field's encoding, this file would differ.
  sp::SessionRequest request;
  if (request.ParseFromString(session_bytes)) {
    std::string reserialized;
    if (request.SerializeToString(&reserialized)) {
      std::ofstream out(output / "cpp_roundtrip.binpb", std::ios::binary);
      out.write(reserialized.data(),
                static_cast<std::streamsize>(reserialized.size()));
    }
  }

  if (g_failures != 0) {
    std::cerr << g_failures << " check(s) failed\n";
    return 1;
  }
  std::cout << "all C++ conformance checks passed\n";
  return 0;
}