klieo_core/

bus.rs

//! Inter-agent bus traits — Pubsub, RequestReply, KvStore, JobQueue.
//!
//! Trait shapes mirror NATS JetStream + KV semantics so the production
//! impl (`klieo-bus-nats`) is a thin adapter, and the in-process impl
//! (`klieo-bus-memory`) can faithfully simulate them. See the spec for
//! reliability invariants.

use crate::error::BusError;
use crate::ids::{DurableName, JobId, RunId};
use async_trait::async_trait;
use bytes::Bytes;
use futures_core::Stream;
use std::collections::HashMap;
use std::pin::Pin;
use std::time::Duration;

/// Opaque headers accompanying a bus message.
pub type Headers = HashMap<String, String>;

/// Bus header carrying the publisher's run id, threaded by
/// `AgentContext::publish` so a receiver can record which run caused a
/// delivery (see [`crate::Episode::BusCausalLink`]).
pub const CAUSATION_HEADER: &str = "klieo-causation-run";

/// One message delivered to a subscriber.
pub struct Msg {
    /// Subject the message was published on.
    pub subject: String,
    /// Payload bytes.
    pub payload: Bytes,
    /// Headers.
    pub headers: Headers,
    /// Acknowledgement handle. The impl provides the underlying mechanism;
    /// callers must invoke exactly one of `ack` / `nak` / `term` per
    /// message or rely on visibility-timeout redelivery.
    pub ack: AckHandle,
}

impl std::fmt::Debug for Msg {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Msg")
            .field("subject", &self.subject)
            .field("payload_len", &self.payload.len())
            .field("headers", &self.headers)
            .finish()
    }
}

/// Acknowledgement handle attached to a delivered [`Msg`].
pub struct AckHandle(Box<dyn AckHandleImpl>);

impl AckHandle {
    /// Construct from an impl. Bus implementations call this when delivering messages.
    pub fn new(inner: Box<dyn AckHandleImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of an ack handle. Exposed so impls can supply backend
/// behaviour without callers reaching into private types.
#[async_trait]
pub trait AckHandleImpl: Send + Sync {
    /// Acknowledge successful processing.
    async fn ack(self: Box<Self>) -> Result<(), BusError>;
    /// Negative-ack with optional redelivery delay. Triggers redelivery
    /// after `delay`.
    async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError>;
    /// Terminate the message (do not redeliver).
    async fn term(self: Box<Self>) -> Result<(), BusError>;
}

impl AckHandle {
    /// Acknowledge.
    pub async fn ack(self) -> Result<(), BusError> {
        self.0.ack().await
    }
    /// Negative-ack.
    pub async fn nak(self, delay: Duration) -> Result<(), BusError> {
        self.0.nak(delay).await
    }
    /// Terminate.
    pub async fn term(self) -> Result<(), BusError> {
        self.0.term().await
    }
}

/// Stream of messages delivered to a subscriber.
pub type MsgStream = Pin<Box<dyn Stream<Item = Result<Msg, BusError>> + Send + 'static>>;

/// Pub/sub interface (subject-based, durable consumers).
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{Headers, Pubsub};
/// use bytes::Bytes;
/// let (pubsub, _, _, _) = noop_bus();
/// pubsub.publish("subject.demo", Bytes::from_static(b"hi"), Headers::new())
///     .await.unwrap();
/// # });
/// ```
#[async_trait]
pub trait Pubsub: Send + Sync {
    /// Publish `payload` on `subject` with `headers`.
    async fn publish(
        &self,
        subject: &str,
        payload: Bytes,
        headers: Headers,
    ) -> Result<(), BusError>;

    /// Subscribe with a durable consumer name. Multiple calls with the
    /// same `durable` form a competing-consumer group sharing replays.
    async fn subscribe(&self, subject: &str, durable: DurableName) -> Result<MsgStream, BusError>;
}

/// Validate that `token` is safe to embed as a single NATS subject segment.
///
/// NATS reserves `.` (separator), `*` (single-token wildcard), and `>`
/// (greedy wildcard) as subject metacharacters. Whitespace, control
/// characters, and non-ASCII bytes are also rejected so wire-level
/// subjects remain printable and the segment cannot collapse the
/// caller-controlled subject namespace.
///
/// Use this at every site that embeds a caller-influenced identifier
/// (progressToken, task id, stream id, …) into a publish or subscribe
/// subject. Skipping validation is a cross-tenant data-leak vector
/// (CWE-74 subject injection): a single-character token like `>` or
/// `*` would subscribe to every other tenant's stream.
///
/// Returns `BusError::Invalid` on any rejected segment.
pub fn validate_subject_token(token: &str) -> Result<(), BusError> {
    if token.is_empty() {
        return Err(BusError::Invalid("subject segment is empty".into()));
    }
    for byte in token.bytes() {
        let forbidden = matches!(byte, b'.' | b'*' | b'>')
            || byte.is_ascii_whitespace()
            || byte.is_ascii_control()
            || !byte.is_ascii();
        if forbidden {
            return Err(BusError::Invalid(format!(
                "subject segment contains forbidden character (byte 0x{byte:02x})"
            )));
        }
    }
    Ok(())
}

#[cfg(test)]
mod subject_token_tests {
    use super::*;

    #[test]
    fn accepts_uuid_like_token() {
        validate_subject_token("550e8400-e29b-41d4-a716-446655440000").unwrap();
        validate_subject_token("task_42").unwrap();
        validate_subject_token("abc123").unwrap();
    }

    #[test]
    fn rejects_empty() {
        let e = validate_subject_token("").unwrap_err();
        assert!(matches!(e, BusError::Invalid(_)));
    }

    #[test]
    fn rejects_greedy_wildcard() {
        assert!(matches!(
            validate_subject_token(">"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_single_token_wildcard() {
        assert!(matches!(
            validate_subject_token("*"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_dot_separator() {
        assert!(matches!(
            validate_subject_token("a.b"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_whitespace_and_control() {
        assert!(matches!(
            validate_subject_token("a b"),
            Err(BusError::Invalid(_))
        ));
        assert!(matches!(
            validate_subject_token("a\nb"),
            Err(BusError::Invalid(_))
        ));
        assert!(matches!(
            validate_subject_token("a\tb"),
            Err(BusError::Invalid(_))
        ));
    }

    #[test]
    fn rejects_non_ascii() {
        assert!(matches!(
            validate_subject_token("café"),
            Err(BusError::Invalid(_))
        ));
    }
}

/// Synchronous request/response over the bus.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{BusError, RequestReply};
/// use bytes::Bytes;
/// use std::time::Duration;
/// let (_, request_reply, _, _) = noop_bus();
/// let err = request_reply
///     .request("svc.add", Bytes::from_static(b"1"), Duration::from_secs(1))
///     .await
///     .unwrap_err();
/// assert!(matches!(err, BusError::NotFound(_)));
/// # });
/// ```
#[async_trait]
pub trait RequestReply: Send + Sync {
    /// Send a request and await one reply, bounded by `timeout`.
    async fn request(
        &self,
        subject: &str,
        payload: Bytes,
        timeout: Duration,
    ) -> Result<Bytes, BusError>;
}

/// CAS revision returned by KV writes.
pub type Revision = u64;

/// One KV entry.
#[derive(Debug, Clone)]
pub struct KvEntry {
    /// Stored value.
    pub value: Bytes,
    /// Revision number after the last write.
    pub revision: Revision,
}

/// One page of bucket keys plus an opaque cursor for the next page.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct KeyPage {
    /// Keys in this page, ascending. Empty iff no key sorts at or after the
    /// requested cursor.
    pub keys: Vec<String>,
    /// Cursor to pass as the next call's `cursor`; `None` once the final key
    /// has been returned.
    pub next: Option<String>,
}

impl KeyPage {
    /// Build a page from its keys and next-cursor. The constructor lets
    /// out-of-crate backends return a `KeyPage` despite `#[non_exhaustive]`.
    pub fn new(keys: Vec<String>, next: Option<String>) -> Self {
        Self { keys, next }
    }
}

/// Bucket-keyed durable KV store with CAS.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::KvStore;
/// use bytes::Bytes;
/// let (_, _, kv, _) = noop_bus();
/// let rev = kv.put("bucket", "key", Bytes::from_static(b"v")).await.unwrap();
/// assert_eq!(rev, 1);
/// # });
/// ```
#[async_trait]
pub trait KvStore: Send + Sync {
    /// Read a value.
    async fn get(&self, bucket: &str, key: &str) -> Result<Option<KvEntry>, BusError>;

    /// Unconditional write. Returns the new revision.
    async fn put(&self, bucket: &str, key: &str, value: Bytes) -> Result<Revision, BusError>;

    /// Compare-and-set. `expected = None` requires the key to be absent;
    /// `Some(rev)` requires the current revision to equal `rev`. Returns
    /// the new revision on success.
    async fn cas(
        &self,
        bucket: &str,
        key: &str,
        value: Bytes,
        expected: Option<Revision>,
    ) -> Result<Revision, BusError>;

    /// Delete a key.
    async fn delete(&self, bucket: &str, key: &str) -> Result<(), BusError>;

    /// Acquire an exclusive lease over `key` for `ttl`. Implementations
    /// should hold the lease as long as the returned [`Lease`] is live
    /// and call its `heartbeat` to extend.
    async fn lease(&self, bucket: &str, key: &str, ttl: Duration) -> Result<Lease, BusError>;

    /// Enumerate keys under `bucket`. Default impl returns
    /// `BusError::Unsupported` — backends that can scan (in-mem,
    /// NATS JetStream KV) override this. The resume-buffer sweeper
    /// (see `klieo-core::resume`) calls this to walk all buckets;
    /// backends that cannot enumerate skip proactive eviction and
    /// rely on access-time TTL checks.
    async fn keys(&self, bucket: &str) -> Result<Vec<String>, BusError> {
        let _ = bucket;
        Err(BusError::Unsupported(
            "keys() not implemented for this KvStore".into(),
        ))
    }

    /// Return all `(key, value)` pairs in `bucket` in a single logical
    /// operation. The default falls back to [`Self::keys`] + per-key
    /// [`Self::get`] (N+1); override for efficiency.
    async fn scan_bucket(&self, bucket: &str) -> Result<Vec<(String, Bytes)>, BusError> {
        let keys = self.keys(bucket).await?;
        let mut out = Vec::with_capacity(keys.len());
        for k in keys {
            if is_causer_key(&k) {
                continue;
            }
            if let Some(entry) = self.get(bucket, &k).await? {
                out.push((k, entry.value));
            }
        }
        Ok(out)
    }

    /// Enumerate `bucket` keys one ascending page at a time. Pass `cursor =
    /// None` for the first page and the prior page's `next` thereafter; paging
    /// ends when `next` is `None`. A `limit` of 0 is floored to 1 so paging
    /// always advances.
    ///
    /// The cursor is the last key of the prior page and resumes strictly after
    /// it, so pages stay stable across concurrent inserts and deletes (unlike
    /// offset paging). The default fetches and sorts every key per call, so it
    /// does NOT bound a large backend's cost — backends over big buckets must
    /// override; the in-process store overrides for a reference impl.
    async fn keys_paginated(
        &self,
        bucket: &str,
        cursor: Option<String>,
        limit: usize,
    ) -> Result<KeyPage, BusError> {
        let mut keys = self.keys(bucket).await?;
        keys.sort();
        Ok(page_from_sorted(keys, cursor.as_deref(), limit))
    }

    /// Record which `run` caused the value at `key` to be written, stored under
    /// a reserved-prefix sibling key in the SAME bucket (no new bucket — NATS KV
    /// bucket names are charset-restricted). First-writer-wins: the causer is
    /// written create-if-absent, so a later writer never overwrites an existing
    /// causer (a CAS conflict is the expected no-op, not an error). Best-effort
    /// provenance, NOT atomic with the value write. The default is backed by
    /// [`Self::cas`]; impls need not override.
    async fn put_causer(&self, bucket: &str, key: &str, run: RunId) -> Result<(), BusError> {
        match self
            .cas(bucket, &causer_key(key), Bytes::from(run.to_string()), None)
            .await
        {
            Ok(_) => Ok(()),
            Err(BusError::CasConflict { .. }) => Ok(()),
            Err(other) => Err(other),
        }
    }

    /// Read the causer recorded by [`Self::put_causer`]; `None` if no causer was
    /// written or the stored value is not a valid run id.
    async fn causer_of(&self, bucket: &str, key: &str) -> Result<Option<RunId>, BusError> {
        let Some(entry) = self.get(bucket, &causer_key(key)).await? else {
            return Ok(None);
        };
        Ok(std::str::from_utf8(&entry.value)
            .ok()
            .and_then(|s| ulid::Ulid::from_string(s).ok())
            .map(RunId))
    }
}

/// Reserved key-prefix under which [`KvStore::put_causer`] stores a value's
/// causer run id, as a sibling key in the same bucket as the value. Uses only
/// NATS-KV-safe characters (`[-/_=.a-zA-Z0-9]`) so the default impl works on
/// every backend; the unusual prefix makes collision with a real key unlikely.
/// `keys()` enumeration still surfaces these sibling keys — callers walking a
/// causation bucket's raw keys should skip them via [`is_causer_key`]
/// ([`KvStore::scan_bucket`] already filters them out).
const KV_CAUSER_PREFIX: &str = "__klieo_causer__.";

fn causer_key(key: &str) -> String {
    format!("{KV_CAUSER_PREFIX}{key}")
}

/// Whether `key` is a reserved causer sidecar key written by
/// [`KvStore::put_causer`] rather than an application value key. Callers that
/// enumerate a causation bucket's raw keys (e.g. a GC sweep over [`KvStore::keys`])
/// should skip keys for which this returns `true`.
pub fn is_causer_key(key: &str) -> bool {
    key.starts_with(KV_CAUSER_PREFIX)
}

/// Slice an ascending, de-duplicated key list into the page that follows
/// `cursor` (exclusive), returning at most `limit` keys plus the next cursor.
/// `limit` is floored to 1 so a zero never stalls a page-walk. Shared by the
/// `keys_paginated` default and by backend overrides that have already gathered
/// a sorted key list.
pub(crate) fn page_from_sorted(keys: Vec<String>, cursor: Option<&str>, limit: usize) -> KeyPage {
    let limit = limit.max(1);
    let start = match cursor {
        Some(c) => keys.partition_point(|k| k.as_str() <= c),
        None => 0,
    };
    let end = start.saturating_add(limit).min(keys.len());
    let page = keys.get(start..end).unwrap_or(&[]).to_vec();
    let next = if end < keys.len() {
        page.last().cloned()
    } else {
        None
    };
    KeyPage::new(page, next)
}

/// Lease handle. Heartbeat to extend; drop to release.
pub struct Lease(Box<dyn LeaseImpl>);

impl Lease {
    /// Construct from an impl. KvStore implementations call this when granting a lease.
    pub fn new(inner: Box<dyn LeaseImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of a lease.
#[async_trait]
pub trait LeaseImpl: Send + Sync {
    /// Extend the TTL.
    async fn heartbeat(&self) -> Result<(), BusError>;
}

impl Lease {
    /// Extend the TTL.
    pub async fn heartbeat(&self) -> Result<(), BusError> {
        self.0.heartbeat().await
    }
}

/// Job enqueued for durable processing.
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Job {
    /// Job payload.
    pub payload: Bytes,
    /// Optional dedup key — if set, the impl writes a `dedup.<queue>`
    /// idempotency record before invoking the handler.
    pub dedup_key: Option<String>,
    /// Maximum redelivery attempts before routing to the DLQ subject.
    /// `None` = use queue default (5).
    pub max_attempts: Option<u32>,
    /// Run id of the agent that caused this job to be enqueued — threaded by
    /// `AgentContext::enqueue`; `None` for un-attributed enqueues. Surfaced on
    /// `ClaimedJob` so a worker can record a causation link.
    pub causation_run_id: Option<RunId>,
}

impl Job {
    /// Build a job with default settings from raw bytes.
    pub fn new(payload: impl Into<Bytes>) -> Self {
        Self {
            payload: payload.into(),
            dedup_key: None,
            max_attempts: None,
            causation_run_id: None,
        }
    }

    /// Start a fluent builder. Chain `.dedup(k)`, `.max_attempts(n)`, and
    /// `.caused_by(run)` then call `.build()`.
    pub fn builder(payload: impl Into<Bytes>) -> JobBuilder {
        JobBuilder {
            payload: payload.into(),
            dedup_key: None,
            max_attempts: None,
            causation_run_id: None,
        }
    }

    /// Stamp the causing run; overwrites any value already set. Typically
    /// called for you by [`crate::agent::AgentContext::enqueue`].
    pub fn caused_by(mut self, run: RunId) -> Self {
        self.causation_run_id = Some(run);
        self
    }
}

/// Fluent builder for [`Job`].
pub struct JobBuilder {
    payload: Bytes,
    dedup_key: Option<String>,
    max_attempts: Option<u32>,
    causation_run_id: Option<RunId>,
}

impl JobBuilder {
    /// Set the dedup key. The queue impl writes a `dedup.<queue>`
    /// idempotency record before invoking the handler.
    pub fn dedup(mut self, key: impl Into<String>) -> Self {
        self.dedup_key = Some(key.into());
        self
    }

    /// Set the maximum redelivery attempts before routing to the DLQ.
    pub fn max_attempts(mut self, n: u32) -> Self {
        self.max_attempts = Some(n);
        self
    }

    /// Set the run id that caused this job to be enqueued.
    pub fn caused_by(mut self, run: RunId) -> Self {
        self.causation_run_id = Some(run);
        self
    }

    /// Finalise the builder into a [`Job`].
    pub fn build(self) -> Job {
        Job {
            payload: self.payload,
            dedup_key: self.dedup_key,
            max_attempts: self.max_attempts,
            causation_run_id: self.causation_run_id,
        }
    }
}

/// Job claimed by a worker.
#[non_exhaustive]
pub struct ClaimedJob {
    /// Job id.
    pub id: JobId,
    /// Job payload.
    pub payload: Bytes,
    /// Lease handle. Caller heartbeats; on drop without ack/nak/dlq the
    /// lease expires and the impl redelivers.
    pub lease: Lease,
    /// Internal handle the impl uses to mark the claim resolved.
    pub claim: ClaimHandle,
    /// Run id of the agent that caused the originating job to be enqueued,
    /// carried through from [`Job::causation_run_id`]; `None` for
    /// un-attributed enqueues. A worker uses it to record a causation link.
    pub causation_run_id: Option<RunId>,
}

impl std::fmt::Debug for ClaimedJob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ClaimedJob")
            .field("id", &self.id)
            .field("payload_len", &self.payload.len())
            .finish()
    }
}

/// Handle the impl uses to resolve a claim.
pub struct ClaimHandle(Box<dyn ClaimHandleImpl>);

impl ClaimHandle {
    /// Construct from an impl. JobQueue implementations call this when claiming a job.
    pub fn new(inner: Box<dyn ClaimHandleImpl>) -> Self {
        Self(inner)
    }
}

/// Implementor side of a claim handle.
#[async_trait]
pub trait ClaimHandleImpl: Send + Sync {
    /// Successful completion. Releases the lease, marks the job done.
    async fn ack(self: Box<Self>) -> Result<(), BusError>;
    /// Retry with backoff.
    async fn nak(self: Box<Self>, delay: Duration) -> Result<(), BusError>;
    /// Send to DLQ.
    async fn dead_letter(self: Box<Self>, reason: &str) -> Result<(), BusError>;
}

impl ClaimedJob {
    /// Construct a claimed job from its parts. JobQueue implementations call
    /// this when claiming a job; the constructor lets out-of-crate backends
    /// build a `ClaimedJob` despite `#[non_exhaustive]`. `causation_run_id`
    /// defaults to `None`; set it with [`ClaimedJob::with_causation`].
    pub fn new(id: JobId, payload: Bytes, lease: Lease, claim: ClaimHandle) -> Self {
        Self {
            id,
            payload,
            lease,
            claim,
            causation_run_id: None,
        }
    }

    /// Set the causing run carried back to the claimer; `None` clears it.
    pub fn with_causation(mut self, run: Option<RunId>) -> Self {
        self.causation_run_id = run;
        self
    }

    /// Heartbeat the lease.
    pub async fn heartbeat(&self) -> Result<(), BusError> {
        self.lease.heartbeat().await
    }
    /// Acknowledge successful completion.
    pub async fn ack(self) -> Result<(), BusError> {
        self.claim.0.ack().await
    }
    /// Negative-ack; retry after `delay`.
    pub async fn nak(self, delay: Duration) -> Result<(), BusError> {
        self.claim.0.nak(delay).await
    }
    /// Route to DLQ.
    pub async fn dead_letter(self, reason: &str) -> Result<(), BusError> {
        self.claim.0.dead_letter(reason).await
    }
}

/// Durable competing-consumer job queue. **No ordering guarantee** —
/// workloads requiring order must use [`Pubsub`] on a partitioned subject
/// with one consumer per partition instead.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::test_utils::noop_bus;
/// use klieo_core::{Job, JobQueue};
/// use bytes::Bytes;
/// let (_, _, _, jobs) = noop_bus();
/// let id = jobs.enqueue("queue.work", Job::new(Bytes::from_static(b"payload"))).await.unwrap();
/// assert_eq!(id.0, "noop-0");
/// # });
/// ```
#[async_trait]
pub trait JobQueue: Send + Sync {
    /// Enqueue a job. Returns a stable id.
    async fn enqueue(&self, queue: &str, job: Job) -> Result<JobId, BusError>;

    /// Claim the next available job. Returns `None` when the queue is
    /// empty. Implementations may long-poll up to a small bounded
    /// duration before returning `None`.
    async fn claim(
        &self,
        queue: &str,
        worker_id: &str,
        lease_ttl: Duration,
    ) -> Result<Option<ClaimedJob>, BusError>;
}

/// Resolved bundle of bus handles ready to drop into an
/// [`crate::agent::AgentContext`] or an `App`.
///
/// Impl crates (`klieo-bus-memory`, `klieo-bus-nats`) provide `From`
/// conversions; downstream code typically writes
/// `BusHandles::from(MemoryBus::new())` and never destructures the
/// four sub-handles by name.
#[derive(Clone)]
pub struct BusHandles {
    /// Pub/sub.
    pub pubsub: std::sync::Arc<dyn Pubsub>,
    /// KV store.
    pub kv: std::sync::Arc<dyn KvStore>,
    /// Synchronous request/reply.
    pub request_reply: std::sync::Arc<dyn RequestReply>,
    /// Durable job queue.
    pub jobs: std::sync::Arc<dyn JobQueue>,
}

impl BusHandles {
    /// Build directly from four already-`Arc`-wrapped handles.
    /// Most callers go through an impl crate's `From` instead.
    pub fn new(
        pubsub: std::sync::Arc<dyn Pubsub>,
        kv: std::sync::Arc<dyn KvStore>,
        request_reply: std::sync::Arc<dyn RequestReply>,
        jobs: std::sync::Arc<dyn JobQueue>,
    ) -> Self {
        Self {
            pubsub,
            kv,
            request_reply,
            jobs,
        }
    }
}

// ─── W3C tracecontext propagation helpers — cluster 0.23 ──────────────
//
// Bus messages crossing replica boundaries carry W3C tracecontext
// (W3C TR/trace-context) in the standard `traceparent` + `tracestate`
// `Headers` keys so OTEL consumers on the receiving side can stitch
// their spans as children of the publisher's span. No API break —
// `Headers` is already `HashMap<String, String>` on `Pubsub::publish`.
//
// No-op when no global OTEL trace provider is installed: the
// TraceContextPropagator silently emits no headers when the current
// context carries no span. Pre-0.23 deployments that ignore the
// headers continue to work unchanged.
//
// These helpers require the `otel` feature flag. Crates that need
// tracecontext propagation must declare:
//   klieo-core = { ..., features = ["otel"] }

#[cfg(feature = "otel")]
use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
#[cfg(feature = "otel")]
use opentelemetry_sdk::propagation::TraceContextPropagator;

/// Inject the current OpenTelemetry context's tracecontext into
/// `headers` under the standard W3C keys (`traceparent`, optionally
/// `tracestate`).
///
/// Call BEFORE `Pubsub::publish` from inside an active OTEL span
/// (typically reached via `tracing-opentelemetry`'s
/// `OpenTelemetrySpanExt::context` on `Span::current()`). Cluster
/// 0.23's impl crates do the bridging — T1 only supplies the
/// lower-level header injection given an explicit Context.
///
/// Requires feature `otel`.
#[cfg(feature = "otel")]
pub fn inject_traceparent(headers: &mut Headers, context: &opentelemetry::Context) {
    let propagator = TraceContextPropagator::new();
    let mut injector = HeaderMapInjector(headers);
    propagator.inject_context(context, &mut injector);
}

/// Extract a W3C tracecontext from `headers`. Returns the default
/// empty Context when no headers are set; callers should set the
/// result as the parent of subsequent spans only if they care about
/// cross-replica stitching.
///
/// Call AFTER `MsgStream::next` resolves.
///
/// Requires feature `otel`.
#[cfg(feature = "otel")]
pub fn extract_traceparent(headers: &Headers) -> opentelemetry::Context {
    let propagator = TraceContextPropagator::new();
    let extractor = HeaderMapExtractor(headers);
    propagator.extract(&extractor)
}

#[cfg(feature = "otel")]
struct HeaderMapInjector<'a>(&'a mut Headers);

#[cfg(feature = "otel")]
impl<'a> Injector for HeaderMapInjector<'a> {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_string(), value);
    }
}

#[cfg(feature = "otel")]
struct HeaderMapExtractor<'a>(&'a Headers);

#[cfg(feature = "otel")]
impl<'a> Extractor for HeaderMapExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(|s| s.as_str())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|s| s.as_str()).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[allow(dead_code)]
    fn _assert_dyn_pubsub(_: &dyn Pubsub) {}
    #[allow(dead_code)]
    fn _assert_dyn_request(_: &dyn RequestReply) {}
    #[allow(dead_code)]
    fn _assert_dyn_kv(_: &dyn KvStore) {}
    #[allow(dead_code)]
    fn _assert_dyn_jobs(_: &dyn JobQueue) {}

    #[test]
    fn job_builder_defaults_match_job_new() {
        let by_builder = Job::builder(Bytes::from_static(b"x")).build();
        let by_new = Job::new(Bytes::from_static(b"x"));
        assert_eq!(by_builder.payload, by_new.payload);
        assert_eq!(by_builder.dedup_key, by_new.dedup_key);
        assert_eq!(by_builder.max_attempts, by_new.max_attempts);
        assert_eq!(by_builder.causation_run_id, by_new.causation_run_id);
    }

    #[test]
    fn job_carries_causation_run_id() {
        let run = crate::ids::RunId::new();
        let job = Job::new(Bytes::from_static(b"x")).caused_by(run);
        assert_eq!(job.causation_run_id, Some(run));
        let plain = Job::new(Bytes::from_static(b"y"));
        assert_eq!(plain.causation_run_id, None);
        let via_builder = Job::builder(Bytes::from_static(b"z"))
            .caused_by(run)
            .build();
        assert_eq!(via_builder.causation_run_id, Some(run));
    }

    #[test]
    fn job_builder_sets_dedup_and_max_attempts() {
        let job = Job::builder(Bytes::from_static(b"payload"))
            .dedup("idempotency-key-42")
            .max_attempts(7)
            .build();
        assert_eq!(job.payload, Bytes::from_static(b"payload"));
        assert_eq!(job.dedup_key.as_deref(), Some("idempotency-key-42"));
        assert_eq!(job.max_attempts, Some(7));
    }

    #[test]
    fn job_builder_accepts_string_dedup_via_into() {
        let owned = String::from("k");
        let job = Job::builder(Bytes::from_static(b"x")).dedup(owned).build();
        assert_eq!(job.dedup_key.as_deref(), Some("k"));
    }

    #[test]
    fn job_field_assignment_still_compiles() {
        let mut job = Job::new(Bytes::from_static(b"x"));
        job.dedup_key = Some("k".into());
        job.max_attempts = Some(3);
        assert_eq!(job.dedup_key.as_deref(), Some("k"));
        assert_eq!(job.max_attempts, Some(3));
    }

    #[test]
    fn bus_error_unsupported_renders_message() {
        let e = BusError::Unsupported("keys() not implemented".into());
        assert_eq!(
            e.to_string(),
            "unsupported operation: keys() not implemented"
        );
    }

    #[cfg(feature = "otel")]
    #[test]
    fn tracecontext_inject_then_extract_roundtrip() {
        use opentelemetry::trace::{
            SpanContext, SpanId, TraceContextExt, TraceFlags, TraceId, TraceState,
        };

        let trace_id = TraceId::from_hex("0123456789abcdef0123456789abcdef").unwrap();
        let span_id = SpanId::from_hex("0123456789abcdef").unwrap();
        let span_ctx = SpanContext::new(
            trace_id,
            span_id,
            TraceFlags::SAMPLED,
            true,
            TraceState::default(),
        );
        let cx = opentelemetry::Context::new().with_remote_span_context(span_ctx);

        let mut headers: Headers = HashMap::new();
        inject_traceparent(&mut headers, &cx);
        assert!(
            headers.contains_key("traceparent"),
            "traceparent header must be injected"
        );

        let extracted = extract_traceparent(&headers);
        let extracted_span_ctx = extracted.span().span_context().clone();
        assert_eq!(extracted_span_ctx.trace_id(), trace_id);
        assert_eq!(extracted_span_ctx.span_id(), span_id);
    }

    fn owned(keys: &[&str]) -> Vec<String> {
        keys.iter().map(|k| k.to_string()).collect()
    }

    #[test]
    fn page_from_sorted_walks_every_key_once_in_order() {
        let keys = owned(&["a", "b", "c", "d", "e"]);
        let mut seen = Vec::new();
        let mut cursor: Option<String> = None;
        loop {
            let page = page_from_sorted(keys.clone(), cursor.as_deref(), 2);
            assert!(page.keys.len() <= 2, "page never exceeds the limit");
            seen.extend(page.keys);
            match page.next {
                Some(c) => cursor = Some(c),
                None => break,
            }
        }
        assert_eq!(seen, owned(&["a", "b", "c", "d", "e"]));
    }

    #[test]
    fn page_from_sorted_cursor_resumes_strictly_after() {
        // Cursor "b" must skip a and b, even if b was since deleted.
        let page = page_from_sorted(owned(&["a", "b", "c", "d"]), Some("b"), 10);
        assert_eq!(page.keys, owned(&["c", "d"]));
        assert!(page.next.is_none());
    }

    #[test]
    fn page_from_sorted_limit_at_or_past_count_is_terminal() {
        let page = page_from_sorted(owned(&["a", "b"]), None, 2);
        assert_eq!(page.keys, owned(&["a", "b"]));
        assert!(page.next.is_none(), "exactly-fits page has no next cursor");
    }

    #[test]
    fn page_from_sorted_floors_zero_limit_to_one() {
        // limit 0 must not stall the walk: it advances by one key, not zero.
        let page = page_from_sorted(owned(&["a", "b", "c"]), None, 0);
        assert_eq!(page.keys, owned(&["a"]));
        assert_eq!(page.next.as_deref(), Some("a"));
    }

    #[test]
    fn page_from_sorted_empty_and_exhausted_cursor_yield_terminal_pages() {
        let empty = page_from_sorted(Vec::new(), None, 4);
        assert!(empty.keys.is_empty() && empty.next.is_none());

        let past_end = page_from_sorted(owned(&["a", "b"]), Some("z"), 4);
        assert!(past_end.keys.is_empty() && past_end.next.is_none());
    }
}