#include "shard_service.h"

#include <algorithm>
#include <chrono>
#include <utility>

namespace meshnet::worker {

namespace {

void FillWorkerFingerprint(sp::Fingerprint* fp, const WorkerIdentity& identity) {
  fp->set_model_artifact_digest(identity.artifact_digest);
  fp->set_runtime_recipe_digest(identity.recipe_digest);
  fp->set_recipe_id(identity.recipe_id);
  fp->set_recipe_version(identity.recipe_version);
  fp->set_catalogue_version(identity.catalogue_version);
}

void FillWorkerShardRange(sp::ShardRange* range, const WorkerIdentity& identity) {
  range->set_start_layer(identity.start_layer);
  range->set_end_layer(identity.end_layer);
  range->set_effective_start_layer(identity.start_layer);
}

// Strictest-of-both bound: the smallest positive of `a`/`b`, or `fallback` when
// neither is set. Mirrors the `_min` helper in `native_protocol/codec.py`.
uint64_t MinPositive(uint64_t a, uint64_t b, uint64_t fallback) {
  if (a > 0 && b > 0) return std::min(a, b);
  if (a > 0) return a;
  if (b > 0) return b;
  return fallback;
}

int64_t NowUnixNanos() {
  return std::chrono::duration_cast<std::chrono::nanoseconds>(
             std::chrono::system_clock::now().time_since_epoch())
      .count();
}

// Build the standard fail response (a terminal-or-not ShardStatus).
sp::SessionResponse MakeFail(const std::string& route_session_id, const std::string& work_id,
                             uint64_t step, sp::ErrorCode code, const std::string& detail,
                             bool terminal, bool retryable) {
  sp::SessionResponse response;
  sp::ShardStatus* status = response.mutable_status();
  status->set_work_id(work_id);
  status->set_route_session_id(route_session_id);
  status->set_idempotency_step(step);
  status->set_terminal(terminal);
  sp::ShardError* error = status->mutable_error();
  error->set_code(code);
  error->set_detail(detail);
  error->set_retryable(retryable);
  return response;
}

sp::SessionResponse MakeAck(const std::string& work_id, uint64_t step, bool duplicate) {
  sp::SessionResponse response;
  sp::Ack* ack = response.mutable_ack();
  ack->set_work_id(work_id);
  ack->set_idempotency_step(step);
  ack->set_duplicate(duplicate);
  return response;
}

void FillDefaultFlow(sp::FlowControl* fc, const FlowLimits& limits) {
  fc->set_credits_granted(limits.credits_granted);
  fc->set_max_inflight_chunks(limits.max_inflight_chunks);
  fc->set_max_chunk_bytes(limits.max_chunk_bytes);
  fc->set_max_prefill_chunk_tokens(limits.max_prefill_chunk_tokens);
}

}  // namespace

grpc::Status ShardRuntimeServiceImpl::GetCapability(grpc::ServerContext*,
                                                    const sp::CapabilityRequest*,
                                                    sp::CapabilityReport* response) {
  response->set_schema_version(sp::SCHEMA_VERSION_1);
  const WorkerIdentity& identity = engine_.identity();
  const EngineHealth health = engine_.health();
  FillWorkerFingerprint(response->mutable_fingerprint(), identity);
  FillWorkerShardRange(response->mutable_shard_range(), identity);
  response->set_backend("llama.cpp");
  response->set_device("cpu");
  response->set_validated(health.serving);
  response->set_detail(health.detail);
  response->set_max_concurrent_sessions(8);
  response->set_max_context_tokens(131072);
  FillDefaultFlow(response->mutable_flow_control(), limits_);
  response->add_accepted_compression(sp::COMPRESSION_NONE);
  response->add_supported_schema_versions(sp::SCHEMA_VERSION_1);
  response->set_validated_at_unix_nanos(0);
  return grpc::Status::OK;
}

grpc::Status ShardRuntimeServiceImpl::Health(grpc::ServerContext*, const sp::HealthRequest*,
                                             sp::HealthReport* response) {
  response->set_schema_version(sp::SCHEMA_VERSION_1);
  const EngineHealth engine_health = engine_.health();
  response->set_state(engine_health.serving ? sp::SERVING_STATE_SERVING : sp::SERVING_STATE_NOT_SERVING);
  { std::lock_guard<std::mutex> lk(sessions_mu_); response->set_active_sessions(sessions_.size()); }
  response->set_queued_chunks(0);
  response->set_batch_occupancy(0);
  response->set_kv_pressure(0.0f);
  response->set_resident_bytes(engine_health.resident_bytes);
  response->set_detail(engine_health.detail + "; loaded=" + engine_.identity().artifact_digest +
                       " range=[" + std::to_string(engine_.identity().start_layer) + "," +
                       std::to_string(engine_.identity().end_layer) + ")");
  return grpc::Status::OK;
}

FlowLimits ShardRuntimeServiceImpl::NegotiateFlow(const sp::FlowControl& proposed) const {
  FlowLimits out;
  out.max_inflight_chunks = static_cast<uint32_t>(MinPositive(
      proposed.max_inflight_chunks(), limits_.max_inflight_chunks, limits_.max_inflight_chunks));
  const uint64_t credits = MinPositive(proposed.credits_granted(), limits_.credits_granted,
                                       limits_.credits_granted);
  out.credits_granted =
      static_cast<uint32_t>(std::min<uint64_t>(credits, out.max_inflight_chunks));
  out.max_chunk_bytes =
      MinPositive(proposed.max_chunk_bytes(), limits_.max_chunk_bytes, limits_.max_chunk_bytes);
  out.max_prefill_chunk_tokens = static_cast<uint32_t>(MinPositive(
      proposed.max_prefill_chunk_tokens(), limits_.max_prefill_chunk_tokens,
      limits_.max_prefill_chunk_tokens));
  return out;
}

uint32_t ShardRuntimeServiceImpl::MarkCancelled(const std::string& route_session_id,
                                                const std::string& work_id) {
  std::lock_guard<std::mutex> lk(sessions_mu_);
  SessionState& state = sessions_[route_session_id];  // creates on first cancel-before-open
  if (state.max_inflight == 0) {
    // Freshly created placeholder for a Cancel that raced ahead of Open.
    state.credits = limits_.credits_granted;
    state.max_inflight = limits_.max_inflight_chunks;
    state.max_chunk_bytes = limits_.max_chunk_bytes;
  }
  if (work_id.empty()) {
    const bool already = state.cancelled_session;
    state.cancelled_session = true;
    return already ? 0 : 1;
  }
  const bool already = state.cancelled_work.count(work_id) != 0;
  state.cancelled_work.insert(work_id);
  return already ? 0 : 1;
}

grpc::Status ShardRuntimeServiceImpl::Session(
    grpc::ServerContext*,
    grpc::ServerReaderWriter<sp::SessionResponse, sp::SessionRequest>* stream) {
  std::string route_session_id;
  sp::SessionRequest request;

  while (stream->Read(&request)) {
    switch (request.kind_case()) {
      case sp::SessionRequest::kOpen: {
        const sp::SessionOpen& open = request.open();
        route_session_id = open.route_session_id();

        // Reject an incompatible peer at open rather than mid-generation. The
        // worker validates the caller's schema, artifact/recipe identity and
        // requested layer range against its own — it never adopts the caller's
        // claimed identity.
        auto reject_open = [&](sp::ErrorCode code, const std::string& detail) {
          stream->Write(MakeFail(route_session_id, /*work_id=*/"", /*step=*/0, code, detail,
                                 /*terminal=*/true, /*retryable=*/false));
        };
        if (open.schema_version() != sp::SCHEMA_VERSION_1) {
          reject_open(sp::ERROR_CODE_SCHEMA_UNSUPPORTED,
                      "worker serves schema version 1 only");
          return grpc::Status::OK;
        }
        const sp::Fingerprint& fp = open.fingerprint();
        if ((!fp.model_artifact_digest().empty() &&
             fp.model_artifact_digest() != engine_.identity().artifact_digest) ||
            (!fp.runtime_recipe_digest().empty() &&
             fp.runtime_recipe_digest() != engine_.identity().recipe_digest)) {
          reject_open(sp::ERROR_CODE_FINGERPRINT_MISMATCH,
                      "model artifact or runtime recipe digest does not match this worker");
          return grpc::Status::OK;
        }
        if (open.has_shard_range()) {
          const sp::ShardRange& r = open.shard_range();
          const bool within = r.start_layer() == engine_.identity().start_layer &&
                              r.end_layer() == engine_.identity().end_layer &&
                              r.effective_start_layer() == engine_.identity().start_layer;
          if (!within) {
            reject_open(sp::ERROR_CODE_SHARD_RANGE_MISMATCH,
                        "requested layer range is not served by this worker");
            return grpc::Status::OK;
          }
        }

        // Settle the flow-control window with strict worker bounds, then keep
        // the negotiated ceilings on the session so every later check enforces
        // exactly what was agreed — not what the peer proposed.
        const FlowLimits negotiated =
            open.has_proposed_flow_control()
                ? NegotiateFlow(open.proposed_flow_control())
                : limits_;
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          SessionState state;
          state.epoch = open.route_epoch();
          state.credits = negotiated.credits_granted;
          state.max_inflight = negotiated.max_inflight_chunks;
          state.max_chunk_bytes = negotiated.max_chunk_bytes;
          state.max_prefill_chunk_tokens = negotiated.max_prefill_chunk_tokens;
          state.opened = true;
          auto it = sessions_.find(route_session_id);
          if (it != sessions_.end()) {
            // A prior out-of-band Cancel may have marked this session cancelled
            // before Open arrived; preserve that so the work still fails closed.
            state.cancelled_session = it->second.cancelled_session;
            state.cancelled_work = it->second.cancelled_work;
          }
          sessions_[route_session_id] = std::move(state);
        }
        sp::SessionResponse response;
        sp::SessionAccepted* accepted = response.mutable_accepted();
        accepted->set_schema_version(sp::SCHEMA_VERSION_1);
        accepted->set_route_session_id(open.route_session_id());
        accepted->set_route_epoch(open.route_epoch());
        FillDefaultFlow(accepted->mutable_flow_control(), negotiated);
        if (open.accepted_compression_size() > 0) {
          for (int c : open.accepted_compression()) {
            accepted->add_accepted_compression(static_cast<sp::Compression>(c));
          }
        } else {
          accepted->add_accepted_compression(sp::COMPRESSION_NONE);
        }
        // Report the fingerprint the worker actually serves, so a mismatch is
        // visible at open — never a copy of the caller's claimed identity.
        FillWorkerFingerprint(accepted->mutable_fingerprint(), engine_.identity());
        stream->Write(response);
        break;
      }

      case sp::SessionRequest::kChunk: {
        const sp::ActivationChunk& chunk = request.chunk();
        const sp::Envelope& envelope = chunk.envelope();
        const std::string work_id = envelope.work_id();
        const uint64_t step = envelope.idempotency_step();

        // Compute the response under the lock, then write it *after* releasing —
        // holding the lock across a (possibly blocking) Write would deadlock an
        // out-of-band Cancel RPC that needs the same lock.
        sp::SessionResponse response;
        bool terminate = false;
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          auto it = sessions_.find(route_session_id);
          SessionState* state = it != sessions_.end() ? &it->second : nullptr;

          if (state == nullptr || !state->opened) {
            // Fail closed: an activation before a valid SessionOpen must never
            // bypass lifecycle, cancellation, epoch or flow-control state.
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
                               "activation received before SessionOpen", true, false);
            terminate = true;
          } else if (state->cancelled_session || state->cancelled_work.count(work_id)) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CANCELLED,
                               "work was cancelled", false, false);
          } else if (envelope.route_epoch() < state->epoch) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_EPOCH_STALE,
                               "stale route epoch", false, false);
          } else if (envelope.deadline_unix_nanos() != 0 &&
                     NowUnixNanos() > envelope.deadline_unix_nanos()) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_DEADLINE_EXCEEDED,
                               "deadline already passed", false, false);
          } else if (state->seen_steps.count(step)) {
            response = MakeAck(work_id, step, /*duplicate=*/true);
          } else if (state->credits <= 0) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_FLOW_CONTROL_VIOLATION,
                               "no flow-control credit remaining", false, true);
          } else {
            const BundleCheck check = engine_.Validate(chunk.bundle(), state->max_chunk_bytes);
            if (check.oversize_detail) {
              response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
                                 *check.oversize_detail, false, false);
            } else if (check.corrupt_detail) {
              response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
                                 *check.corrupt_detail, false, false);
            } else {
              state->seen_steps.insert(step);
              state->credits -= 1;
              std::string execution_error;
              if (!engine_.Execute(chunk.bundle(), &execution_error)) {
                response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
                                    execution_error, false, true);
              } else {
              *response.mutable_chunk() = chunk;  // echo the exact bundle back
              }
            }
          }
        }
        stream->Write(response);
        if (terminate) {
          return grpc::Status::OK;
        }
        break;
      }

      case sp::SessionRequest::kDecode: {
        const sp::DecodeStep& step_msg = request.decode();
        const std::string work_id = step_msg.work_id();
        const uint64_t step = step_msg.idempotency_step();

        sp::TensorBundle bundle;
        if (step_msg.bundle().tensors_size() > 0) {
          bundle = step_msg.bundle();
        } else {
          bundle.set_bundle_version(1);
          *bundle.add_tensors() = step_msg.tensor();
        }

        sp::SessionResponse response;
        bool terminate = false;
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          auto it = sessions_.find(route_session_id);
          SessionState* state = it != sessions_.end() ? &it->second : nullptr;

          if (state == nullptr || !state->opened) {
            // Fail closed: a decode step before a valid SessionOpen must never
            // bypass lifecycle, cancellation, epoch or flow-control state.
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
                               "activation received before SessionOpen", true, false);
            terminate = true;
          } else if (state->cancelled_session || state->cancelled_work.count(work_id)) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_CANCELLED,
                               "work was cancelled", false, false);
          } else if (step_msg.deadline_unix_nanos() != 0 &&
                     NowUnixNanos() > step_msg.deadline_unix_nanos()) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_DEADLINE_EXCEEDED,
                               "deadline already passed", false, false);
          } else if (state->seen_steps.count(step)) {
            response = MakeAck(work_id, step, /*duplicate=*/true);
          } else if (state->credits <= 0) {
            response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_FLOW_CONTROL_VIOLATION,
                               "no flow-control credit remaining", false, true);
          } else {
            const BundleCheck check = engine_.Validate(bundle, state->max_chunk_bytes);
            if (check.oversize_detail) {
              response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_RESOURCE_EXHAUSTED,
                                 *check.oversize_detail, false, false);
            } else if (check.corrupt_detail) {
              response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_PAYLOAD_CORRUPT,
                                 *check.corrupt_detail, false, false);
            } else {
              state->seen_steps.insert(step);
              state->credits -= 1;
              std::string execution_error;
              if (!engine_.Execute(bundle, &execution_error)) {
                response = MakeFail(route_session_id, work_id, step, sp::ERROR_CODE_INTERNAL,
                                    execution_error, false, true);
              } else {
              // No decode response field exists; echo the step back as a
              // chunk-bearing SessionResponse per the proto's relayed-frame design.
              sp::ActivationChunk* out = response.mutable_chunk();
              sp::Envelope* out_env = out->mutable_envelope();
              out_env->set_schema_version(sp::SCHEMA_VERSION_1);
              out_env->set_work_id(work_id);
              out_env->set_idempotency_step(step);
              out_env->set_phase(sp::PHASE_DECODE);
              sp::PositionSpan* pos = out_env->mutable_position();
              pos->set_first_position(step_msg.position());
              pos->set_token_count(1);
              *out->mutable_bundle() = bundle;
              }
            }
          }
        }
        stream->Write(response);
        if (terminate) {
          return grpc::Status::OK;
        }
        break;
      }

      case sp::SessionRequest::kFlowControl: {
        const uint32_t topup = request.flow_control().credits_granted();
        sp::SessionResponse response;
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          auto it = sessions_.find(route_session_id);
          sp::FlowControl* fc = response.mutable_flow_control();
          if (it != sessions_.end()) {
            SessionState& state = it->second;
            int64_t granted = std::min<int64_t>(state.credits + topup,
                                                 static_cast<int64_t>(state.max_inflight));
            state.credits = granted;
            fc->set_credits_granted(static_cast<uint32_t>(granted));
            fc->set_max_inflight_chunks(state.max_inflight);
            fc->set_max_chunk_bytes(state.max_chunk_bytes);
          } else {
            fc->set_credits_granted(topup != 0 ? topup : limits_.credits_granted);
            fc->set_max_inflight_chunks(limits_.max_inflight_chunks);
            fc->set_max_chunk_bytes(limits_.max_chunk_bytes);
          }
          fc->set_max_prefill_chunk_tokens(limits_.max_prefill_chunk_tokens);
        }
        stream->Write(response);
        break;
      }

      case sp::SessionRequest::kRelease: {
        const sp::ReleaseSignal& release = request.release();
        // An explicit release drops session state immediately (KV, credits,
        // dedup) instead of holding it for the TTL — the whole point of the
        // signal. Erase the session this stream opened so its resources are
        // freed the moment the terminal status is sent.
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          sessions_.erase(route_session_id);
          engine_.ReleaseSession(route_session_id);
        }
        sp::SessionResponse response;
        sp::ShardStatus* status = response.mutable_status();
        status->set_work_id(release.work_id());
        status->set_route_session_id(release.route_session_id());
        status->set_terminal(true);
        stream->Write(response);
        return grpc::Status::OK;
      }

      case sp::SessionRequest::kCancel: {
        const sp::CancelSignal& signal = request.cancel();
        MarkCancelled(route_session_id, signal.work_id());
        const bool whole_session = signal.work_id().empty();
        stream->Write(MakeFail(route_session_id, signal.work_id(), 0, sp::ERROR_CODE_CANCELLED,
                               signal.reason().empty() ? "cancelled" : signal.reason(),
                               whole_session, false));
        if (whole_session) {
          return grpc::Status::OK;
        }
        break;
      }

      default: {
        sp::SessionResponse response;
        response.mutable_status()->set_terminal(true);
        stream->Write(response);
        return grpc::Status::OK;
      }
    }
  }
  return grpc::Status::OK;
}

grpc::Status ShardRuntimeServiceImpl::Release(grpc::ServerContext*,
                                              const sp::ReleaseRequest* request,
                                              sp::ReleaseResponse* response) {
  bool existed;
  {
    std::lock_guard<std::mutex> lk(sessions_mu_);
    existed = sessions_.erase(request->route_session_id()) != 0;
    engine_.ReleaseSession(request->route_session_id());
  }
  response->set_released(existed);
  return grpc::Status::OK;
}

grpc::Status ShardRuntimeServiceImpl::Cancel(grpc::ServerContext*,
                                             const sp::CancelRequest* request,
                                             sp::CancelResponse* response) {
  const uint32_t newly = MarkCancelled(request->route_session_id(), request->work_id());
  response->set_cancelled_work_items(newly);
  return grpc::Status::OK;
}

}  // namespace meshnet::worker