// The native Shard worker's ShardRuntime service (DGR-033).
//
// A faithful C++ port of `ShardRuntimeServicer` in `shard_runtime_server.py`:
// the same per-`route_session_id` identity/credit/dedup state, the same
// fail-closed negative paths (stale epoch, expired deadline, corrupt/oversize
// payload, exhausted flow-control credit, duplicate idempotency step, in-band
// and out-of-band cancellation), and the same lifecycle (open/prefill/decode/
// flow-control/release/cancel). The only compute it does is the fake engine's
// bounded forward — there is no llama.cpp linkage and no arbitrary-graph RPC.

#ifndef MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_
#define MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_

#include <cstdint>
#include <map>
#include <mutex>
#include <set>
#include <string>

#include <grpcpp/grpcpp.h>

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

namespace meshnet::worker {

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

struct FlowLimits {
  uint32_t credits_granted = 16;
  uint32_t max_inflight_chunks = 16;
  uint64_t max_chunk_bytes = 4u * 1024u * 1024u;
  uint32_t max_prefill_chunk_tokens = 512;
};

// Per-route-session identity/credit/dedup state, kept on the servicer instance
// (guarded by a lock) so an out-of-band unary Cancel from a different handler
// thread can reach a session a concurrent Session stream is still iterating.
struct SessionState {
  uint64_t epoch = 0;
  int64_t credits = 0;
  uint32_t max_inflight = 0;
  uint64_t max_chunk_bytes = 0;
  uint32_t max_prefill_chunk_tokens = 0;
  std::set<uint64_t> seen_steps;
  std::set<std::string> cancelled_work;
  bool cancelled_session = false;
  // True only after a valid SessionOpen handshake completed for this
  // route_session_id. An activation (chunk/decode) that arrives while this is
  // false fails closed: no work may bypass the lifecycle handshake, even when a
  // placeholder state already exists from an out-of-band Cancel that raced Open.
  bool opened = false;
};

class ShardRuntimeServiceImpl final : public sp::ShardRuntime::Service {
 public:
  ShardRuntimeServiceImpl(FlowLimits limits, ShardEngine& engine) : limits_(limits), engine_(engine) {}

  grpc::Status GetCapability(grpc::ServerContext* context,
                             const sp::CapabilityRequest* request,
                             sp::CapabilityReport* response) override;

  grpc::Status Health(grpc::ServerContext* context, const sp::HealthRequest* request,
                      sp::HealthReport* response) override;

  grpc::Status Session(
      grpc::ServerContext* context,
      grpc::ServerReaderWriter<sp::SessionResponse, sp::SessionRequest>* stream) override;

  grpc::Status Release(grpc::ServerContext* context, const sp::ReleaseRequest* request,
                       sp::ReleaseResponse* response) override;

  grpc::Status Cancel(grpc::ServerContext* context, const sp::CancelRequest* request,
                      sp::CancelResponse* response) override;

 private:
  // Returns the number of items newly marked cancelled, creating session state
  // if the Cancel raced ahead of SessionOpen.
  uint32_t MarkCancelled(const std::string& route_session_id, const std::string& work_id);

  // Settle a stream's flow-control window against this worker's own limits: the
  // strictest bound of either peer wins for every field, so a peer can never
  // raise the worker's ceilings by proposing a larger window. Mirrors
  // `negotiate_flow_control` in `native_protocol/codec.py`.
  FlowLimits NegotiateFlow(const sp::FlowControl& proposed) const;

  FlowLimits limits_;
  ShardEngine& engine_;
  std::mutex sessions_mu_;
  std::map<std::string, SessionState> sessions_;
};

}  // namespace meshnet::worker

#endif  // MESHNET_NATIVE_WORKER_SHARD_SERVICE_H_