// C++ round-trip and cross-language compatibility test for the Shard protocol.
//
// Modes (composable):
//   --selftest         serialize a sample message, parse it back, verify fields.
//   --read <path>      parse a fixture serialized by another language; verify the
//                      known fields; tolerate unknown fields (forward compat).
//   --write <path>     serialize the C++ sample so another language can parse it.
//
// Exit code 0 means every requested check passed. The Python test drives this
// binary to prove Python<->C++ wire compatibility in both directions.

#include "shard_runtime.pb.h"

#include <cstdint>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

using namespace meshnet::shard::v1;

namespace {

bool Fail(const std::string& why) {
  std::cerr << "roundtrip_test: FAIL: " << why << std::endl;
  return false;
}

SessionActivation MakeSample() {
  SessionActivation act;
  PrefillChunk* pre = act.mutable_prefill();

  MessageHeader* h = pre->mutable_header();
  h->set_schema_version(SCHEMA_VERSION_1);
  h->set_work_id("w1");
  h->set_route_session_id("s1");
  h->set_route_epoch(3);
  h->set_phase(PHASE_PREFILL);
  h->set_idempotency_step(7);
  h->set_cache_expectation(CACHE_FRESH);
  h->set_compression(COMPRESSION_NONE);

  ArtifactFingerprint* fp = h->mutable_fingerprint();
  fp->set_model_id("meta-llama/Llama-3.1-8B");
  fp->set_quantization("Q4_K_M");
  fp->set_runtime_recipe_fingerprint("recipe-abc");

  ShardRange* sr = h->mutable_shard_range();
  sr->set_start_layer(0);
  sr->set_end_layer(16);
  sr->set_effective_start_layer(0);
  sr->set_owns_embedding(true);

  Position* pos = h->mutable_position();
  pos->set_start_position(0);
  pos->set_token_count(5);
  pos->set_sequence_length(5);

  pre->set_chunk_index(0);
  pre->set_chunk_count(1);
  pre->set_final_chunk(true);

  TensorBundle* bundle = pre->mutable_activations();
  bundle->set_bundle_version(1);
  NamedTensor* t = bundle->add_tensors();
  t->set_name("hidden");
  t->add_shape(1);
  t->add_shape(4096);
  t->set_dtype(DTYPE_F16);
  t->set_byte_order(BYTE_ORDER_LITTLE_ENDIAN);
  t->set_total_byte_length(8);
  t->set_compression(COMPRESSION_NONE);
  TensorFragment* frag = t->add_fragments();
  frag->set_fragment_index(0);
  frag->set_fragment_count(1);
  frag->set_byte_offset(0);
  frag->set_data(std::string("\x01\x02\x03\x04\x05\x06\x07\x08", 8));

  return act;
}

bool CheckSample(const SessionActivation& act) {
  if (act.payload_case() != SessionActivation::kPrefill)
    return Fail("payload is not prefill");
  const PrefillChunk& pre = act.prefill();
  const MessageHeader& h = pre.header();
  if (h.schema_version() != SCHEMA_VERSION_1) return Fail("schema_version");
  if (h.work_id() != "w1") return Fail("work_id");
  if (h.route_session_id() != "s1") return Fail("route_session_id");
  if (h.route_epoch() != 3) return Fail("route_epoch");
  if (h.phase() != PHASE_PREFILL) return Fail("phase");
  if (h.idempotency_step() != 7) return Fail("idempotency_step");
  if (h.fingerprint().model_id() != "meta-llama/Llama-3.1-8B")
    return Fail("model_id");
  if (h.fingerprint().quantization() != "Q4_K_M") return Fail("quantization");
  if (h.shard_range().end_layer() != 16) return Fail("end_layer");
  if (!h.shard_range().owns_embedding()) return Fail("owns_embedding");
  if (h.position().token_count() != 5) return Fail("token_count");
  if (!pre.final_chunk()) return Fail("final_chunk");
  if (pre.activations().tensors_size() != 1) return Fail("tensors_size");
  const NamedTensor& t = pre.activations().tensors(0);
  if (t.name() != "hidden") return Fail("tensor name");
  if (t.dtype() != DTYPE_F16) return Fail("dtype");
  if (t.byte_order() != BYTE_ORDER_LITTLE_ENDIAN) return Fail("byte_order");
  if (t.shape_size() != 2 || t.shape(1) != 4096) return Fail("shape");
  if (t.fragments_size() != 1) return Fail("fragments_size");
  if (t.fragments(0).data().size() != 8) return Fail("fragment data length");
  return true;
}

bool ReadFile(const std::string& path, std::string* out) {
  std::ifstream in(path, std::ios::binary);
  if (!in) return false;
  std::ostringstream ss;
  ss << in.rdbuf();
  *out = ss.str();
  return true;
}

bool WriteFile(const std::string& path, const std::string& data) {
  std::ofstream out(path, std::ios::binary);
  if (!out) return false;
  out.write(data.data(), static_cast<std::streamsize>(data.size()));
  return static_cast<bool>(out);
}

}  // namespace

int main(int argc, char** argv) {
  GOOGLE_PROTOBUF_VERIFY_VERSION;

  std::string read_path;
  std::string write_path;
  bool selftest = (argc == 1);

  for (int i = 1; i < argc; ++i) {
    std::string arg = argv[i];
    if (arg == "--selftest") {
      selftest = true;
    } else if (arg == "--read" && i + 1 < argc) {
      read_path = argv[++i];
    } else if (arg == "--write" && i + 1 < argc) {
      write_path = argv[++i];
    } else {
      std::cerr << "unknown/incomplete arg: " << arg << std::endl;
      return 2;
    }
  }

  if (selftest) {
    SessionActivation sample = MakeSample();
    std::string bytes;
    if (!sample.SerializeToString(&bytes)) return Fail("serialize"), 1;
    SessionActivation parsed;
    if (!parsed.ParseFromString(bytes)) return Fail("parse"), 1;
    if (!CheckSample(parsed)) return 1;
    std::cout << "selftest ok (" << bytes.size() << " bytes)" << std::endl;
  }

  if (!read_path.empty()) {
    std::string bytes;
    if (!ReadFile(read_path, &bytes)) return Fail("cannot read fixture"), 1;
    SessionActivation parsed;
    // ParseFromString tolerates and preserves unknown fields (forward compat).
    if (!parsed.ParseFromString(bytes)) return Fail("parse fixture"), 1;
    if (!CheckSample(parsed)) return 1;
    std::cout << "read ok (" << bytes.size() << " bytes)" << std::endl;
  }

  if (!write_path.empty()) {
    SessionActivation sample = MakeSample();
    std::string bytes;
    if (!sample.SerializeToString(&bytes)) return Fail("serialize for write"), 1;
    if (!WriteFile(write_path, bytes)) return Fail("cannot write output"), 1;
    std::cout << "write ok (" << bytes.size() << " bytes)" << std::endl;
  }

  google::protobuf::ShutdownProtobufLibrary();
  return 0;
}