// Standalone native Shard worker executable (DGR-033).
//
// Serves the complete ShardRuntime lifecycle/stream contract over real
// gRPC/HTTP2 using the model-free FakeShardEngine. It links neither llama.cpp
// nor any graph-execution entry point: the only surface it exposes is the
// ShardRuntime service defined in shard_runtime.proto.
//
// Usage:
//   shard_worker [listen_addr]     serve until SIGTERM/SIGINT (graceful drain)
//   shard_worker --selftest        bind an ephemeral port, self-drive the
//                                  lifecycle over a real loopback channel, exit
//
// Environment:
//   MESHNET_SHARD_LISTEN_ADDR   host:port to bind (default localhost:50051)
//   MESHNET_MAX_CHUNK_BYTES     per-chunk byte ceiling the worker enforces
//
// On a normal run it prints one readiness line — "ShardRuntime worker listening
// on <addr>" — once the socket is bound, so a supervisor/harness has a real
// readiness signal instead of a sleep.

#include <atomic>
#include <cerrno>
#include <csignal>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <memory>
#include <string>
#include <thread>
#include <unistd.h>

#include <grpcpp/grpcpp.h>

#include "shard_service.h"
#include "shard_runtime.grpc.pb.h"

namespace {

namespace sp = ::meshnet::shard::v1;

// Self-pipe: the signal handler must stay async-signal-safe, so it only writes
// one byte; a helper thread reads it and performs the (non-signal-safe) server
// Shutdown(). Set once in main() before installing the handler.
volatile std::sig_atomic_t g_signal_pipe_write_fd = -1;

extern "C" void HandleTermination(int /*signum*/) {
  if (g_signal_pipe_write_fd >= 0) {
    const char byte = 1;
    ssize_t rc = ::write(g_signal_pipe_write_fd, &byte, 1);
    (void)rc;  // best-effort; nothing safe to do on failure inside a handler
  }
}

meshnet::worker::FlowLimits LimitsFromEnv() {
  meshnet::worker::FlowLimits limits;
  if (const char* raw = std::getenv("MESHNET_MAX_CHUNK_BYTES")) {
    char* end = nullptr;
    const unsigned long long value = std::strtoull(raw, &end, 10);
    if (end != raw && value > 0) {
      limits.max_chunk_bytes = static_cast<uint64_t>(value);
    }
  }
  return limits;
}

int RunSelfTest() {
  meshnet::worker::ShardRuntimeServiceImpl service(LimitsFromEnv());
  int selected_port = 0;
  grpc::ServerBuilder builder;
  builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &selected_port);
  builder.RegisterService(&service);
  std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
  if (!server || selected_port == 0) {
    std::cerr << "selftest: failed to bind ephemeral port\n";
    return 1;
  }
  const std::string target = "127.0.0.1:" + std::to_string(selected_port);
  auto channel = grpc::CreateChannel(target, grpc::InsecureChannelCredentials());
  auto stub = sp::ShardRuntime::NewStub(channel);

  int failures = 0;
  auto check = [&](bool cond, const char* what) {
    if (!cond) {
      std::cerr << "selftest FAIL: " << what << "\n";
      ++failures;
    }
  };

  // Capability + health.
  {
    grpc::ClientContext ctx;
    sp::CapabilityRequest req;
    req.set_schema_version(sp::SCHEMA_VERSION_1);
    sp::CapabilityReport rep;
    grpc::Status status = stub->GetCapability(&ctx, req, &rep);
    check(status.ok(), "GetCapability RPC");
    check(rep.validated(), "capability validated");
    check(rep.schema_version() == sp::SCHEMA_VERSION_1, "capability schema version");
  }
  {
    grpc::ClientContext ctx;
    sp::HealthRequest req;
    req.set_schema_version(sp::SCHEMA_VERSION_1);
    sp::HealthReport rep;
    grpc::Status status = stub->Health(&ctx, req, &rep);
    check(status.ok(), "Health RPC");
    check(rep.state() == sp::SERVING_STATE_SERVING, "health serving");
  }

  // A minimal session: open -> fragmented prefill -> decode -> release.
  {
    grpc::ClientContext ctx;
    auto stream = stub->Session(&ctx);

    sp::SessionRequest open;
    sp::SessionOpen* o = open.mutable_open();
    o->set_schema_version(sp::SCHEMA_VERSION_1);
    o->set_route_session_id("selftest");
    o->set_route_epoch(1);
    sp::FlowControl* fc = o->mutable_proposed_flow_control();
    fc->set_credits_granted(16);
    fc->set_max_inflight_chunks(16);
    fc->set_max_chunk_bytes(4u * 1024u * 1024u);
    check(stream->Write(open), "write open");

    sp::SessionResponse accepted;
    check(stream->Read(&accepted), "read accepted");
    check(accepted.kind_case() == sp::SessionResponse::kAccepted, "accepted kind");

    // Fragmented prefill: two fragments tiling a 6-byte payload.
    const std::string payload = "ABCDEF";
    sp::SessionRequest chunk;
    sp::ActivationChunk* ac = chunk.mutable_chunk();
    sp::Envelope* env = ac->mutable_envelope();
    env->set_schema_version(sp::SCHEMA_VERSION_1);
    env->set_work_id("w1");
    env->set_route_session_id("selftest");
    env->set_route_epoch(1);
    env->set_idempotency_step(1);
    env->set_phase(sp::PHASE_PREFILL);
    sp::TensorBundle* bundle = ac->mutable_bundle();
    bundle->set_bundle_version(1);
    sp::NamedTensor* tensor = bundle->add_tensors();
    tensor->set_name("hidden_states");
    tensor->set_dtype(sp::DTYPE_BFLOAT16);
    tensor->set_byte_order(sp::BYTE_ORDER_LITTLE_ENDIAN);
    tensor->set_total_bytes(payload.size());
    tensor->set_compression(sp::COMPRESSION_NONE);
    sp::Checksum* cksum = tensor->mutable_checksum();
    cksum->set_algorithm(sp::CHECKSUM_ALGORITHM_CRC32C);
    const uint32_t crc = meshnet::worker::Crc32(payload);
    std::string crc_be(4, '\0');
    crc_be[0] = static_cast<char>((crc >> 24) & 0xFF);
    crc_be[1] = static_cast<char>((crc >> 16) & 0xFF);
    crc_be[2] = static_cast<char>((crc >> 8) & 0xFF);
    crc_be[3] = static_cast<char>(crc & 0xFF);
    cksum->set_value(crc_be);
    sp::TensorFragment* f0 = tensor->add_fragments();
    f0->set_fragment_index(0);
    f0->set_fragment_count(2);
    f0->set_byte_offset(0);
    f0->set_payload(payload.substr(0, 3));
    sp::TensorFragment* f1 = tensor->add_fragments();
    f1->set_fragment_index(1);
    f1->set_fragment_count(2);
    f1->set_byte_offset(3);
    f1->set_payload(payload.substr(3));
    check(stream->Write(chunk), "write chunk");

    sp::SessionResponse echoed;
    check(stream->Read(&echoed), "read chunk echo");
    check(echoed.kind_case() == sp::SessionResponse::kChunk, "chunk echo kind");

    sp::SessionRequest decode;
    sp::DecodeStep* ds = decode.mutable_decode();
    ds->set_idempotency_step(2);
    ds->set_position(1);
    ds->set_work_id("w2");
    sp::TensorBundle* dbundle = ds->mutable_bundle();
    dbundle->set_bundle_version(1);
    sp::NamedTensor* dt = dbundle->add_tensors();
    dt->set_name("hidden_states");
    dt->set_dtype(sp::DTYPE_BFLOAT16);
    dt->set_byte_order(sp::BYTE_ORDER_LITTLE_ENDIAN);
    dt->set_total_bytes(payload.size());
    dt->set_compression(sp::COMPRESSION_NONE);
    sp::Checksum* dck = dt->mutable_checksum();
    dck->set_algorithm(sp::CHECKSUM_ALGORITHM_CRC32C);
    dck->set_value(crc_be);
    sp::TensorFragment* df = dt->add_fragments();
    df->set_fragment_index(0);
    df->set_fragment_count(1);
    df->set_byte_offset(0);
    df->set_payload(payload);
    check(stream->Write(decode), "write decode");

    sp::SessionResponse decode_echo;
    check(stream->Read(&decode_echo), "read decode echo");
    check(decode_echo.kind_case() == sp::SessionResponse::kChunk, "decode echo kind");

    sp::SessionRequest release;
    sp::ReleaseSignal* rs = release.mutable_release();
    rs->set_route_session_id("selftest");
    rs->set_work_id("w-final");
    check(stream->Write(release), "write release");
    stream->WritesDone();

    sp::SessionResponse terminal;
    check(stream->Read(&terminal), "read terminal");
    check(terminal.kind_case() == sp::SessionResponse::kStatus && terminal.status().terminal(),
          "terminal status");

    grpc::Status status = stream->Finish();
    check(status.ok(), "stream finish");
  }

  server->Shutdown();
  server->Wait();

  if (failures == 0) {
    std::cout << "selftest: all lifecycle checks passed\n";
    return 0;
  }
  std::cerr << "selftest: " << failures << " check(s) failed\n";
  return 1;
}

}  // namespace

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

  for (int i = 1; i < argc; ++i) {
    if (std::strcmp(argv[i], "--selftest") == 0) {
      return RunSelfTest();
    }
  }

  std::string listen_addr = "localhost:50051";
  if (const char* env = std::getenv("MESHNET_SHARD_LISTEN_ADDR")) {
    listen_addr = env;
  }
  if (argc > 1 && argv[1][0] != '-') {
    listen_addr = argv[1];
  }

  meshnet::worker::FlowLimits limits = LimitsFromEnv();
  meshnet::worker::ShardRuntimeServiceImpl service(limits);

  grpc::ServerBuilder builder;
  int selected_port = 0;
  builder.AddListeningPort(listen_addr, grpc::InsecureServerCredentials(), &selected_port);
  // Bounded messages, two layers: a hard transport receive ceiling (never below
  // 4 MiB so the handshake and normal chunks always fit) plus the finer
  // app-level per-tensor RESOURCE_EXHAUSTED check the service enforces against
  // the negotiated max_chunk_bytes. Neither path lets an unbounded frame in.
  constexpr int kTransportFloor = 4 * 1024 * 1024;
  const int transport_max = limits.max_chunk_bytes > static_cast<uint64_t>(kTransportFloor)
                                ? static_cast<int>(limits.max_chunk_bytes)
                                : kTransportFloor;
  builder.SetMaxReceiveMessageSize(transport_max);
  builder.RegisterService(&service);
  std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
  if (!server || selected_port == 0) {
    std::cerr << "failed to bind " << listen_addr << "\n";
    return 1;
  }

  int pipe_fds[2];
  if (::pipe(pipe_fds) != 0) {
    std::cerr << "failed to create shutdown pipe\n";
    return 1;
  }
  g_signal_pipe_write_fd = pipe_fds[1];

  struct sigaction sa;
  std::memset(&sa, 0, sizeof(sa));
  sa.sa_handler = HandleTermination;
  ::sigaction(SIGTERM, &sa, nullptr);
  ::sigaction(SIGINT, &sa, nullptr);

  // Drain thread: wakes on the first termination signal and shuts the server
  // down gracefully so in-flight sessions finish rather than being severed.
  std::thread drain([&server, read_fd = pipe_fds[0]]() {
    char byte = 0;
    ssize_t rc = 0;
    do {
      rc = ::read(read_fd, &byte, 1);
    } while (rc < 0 && errno == EINTR);
    server->Shutdown();
  });

  std::cout << "ShardRuntime worker listening on " << listen_addr << std::endl;

  server->Wait();
  drain.join();
  ::close(pipe_fds[0]);
  ::close(pipe_fds[1]);
  std::cout << "ShardRuntime worker shut down cleanly" << std::endl;
  return 0;
}