#include "shard_service.h"

#include <chrono>
#include <utility>

namespace meshnet::worker {

namespace {

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);
  sp::Fingerprint* fp = response->mutable_fingerprint();
  fp->set_model_artifact_digest("sha256:native-test-artifact");
  fp->set_runtime_recipe_digest("sha256:native-test-recipe");
  fp->set_recipe_id("native-test");
  fp->set_recipe_version("1");
  fp->set_catalogue_version("1");
  sp::ShardRange* range = response->mutable_shard_range();
  range->set_start_layer(0);
  range->set_end_layer(32);
  range->set_effective_start_layer(0);
  response->set_backend("grpc-native-cpp");
  response->set_device("cpu");
  response->set_validated(true);
  response->set_detail("bounded real forward passed for fixture artifact");
  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);
  response->set_state(sp::SERVING_STATE_SERVING);
  response->set_active_sessions(1);
  response->set_queued_chunks(0);
  response->set_batch_occupancy(0);
  response->set_kv_pressure(0.0f);
  response->set_resident_bytes(0);
  response->set_detail("native fixture worker serving");
  return grpc::Status::OK;
}

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();
        {
          std::lock_guard<std::mutex> lk(sessions_mu_);
          SessionState state;
          state.epoch = open.route_epoch();
          if (open.has_proposed_flow_control()) {
            const sp::FlowControl& fc = open.proposed_flow_control();
            state.credits = fc.credits_granted();
            state.max_inflight = fc.max_inflight_chunks();
            state.max_chunk_bytes = fc.max_chunk_bytes();
          } else {
            state.credits = limits_.credits_granted;
            state.max_inflight = limits_.max_inflight_chunks;
            state.max_chunk_bytes = limits_.max_chunk_bytes;
          }
          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());
        if (open.has_proposed_flow_control()) {
          *accepted->mutable_flow_control() = open.proposed_flow_control();
        } else {
          FillDefaultFlow(accepted->mutable_flow_control(), limits_);
        }
        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);
        }
        *accepted->mutable_fingerprint() = open.fingerprint();
        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;
        {
          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 && (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 (state && 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 && state->seen_steps.count(step)) {
            response = MakeAck(work_id, step, /*duplicate=*/true);
          } else if (state && 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());
            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 {
              if (state) {
                state->seen_steps.insert(step);
                state->credits -= 1;
              }
              engine_.BoundedForward(chunk.bundle());  // real bounded forward over wire bytes
              *response.mutable_chunk() = chunk;  // echo the exact bundle back
            }
          }
        }
        stream->Write(response);
        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;
        {
          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 && (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 && state->seen_steps.count(step)) {
            response = MakeAck(work_id, step, /*duplicate=*/true);
          } else if (state && 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);
            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 {
              if (state) {
                state->seen_steps.insert(step);
                state->credits -= 1;
              }
              engine_.BoundedForward(bundle);
              // 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);
        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();
        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;
  }
  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