//! W3C Trace Context (RFC 9462) types + `_meta` propagation
//! (SEP-414).
//!
//! The 2026-07-28 spec mandates that an MCP client propagating a
//! parent span ride the W3C Trace Context format through the request
//! `_meta` block:
//!
//! ```json
//! {
//!   "_meta": {
//!     "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
//!     "tracestate":  "vendor1=opaqueValue,vendor2=opaqueValue"
//!   }
//! }
//! ```
//!
//! v1 ships the typed model + inject / extract. The library is
//! deliberately lossless on the `tracestate` payload (opaque per the
//! spec) and strict on the `traceparent` shape (everything is
//! length-checked and hex-validated, mismatches fail the parse).

use serde_json::Value;
use thiserror::Error;

/// Errors raised by [`TraceParent::parse`] and the extractors.
#[derive(Debug, Error, PartialEq, Eq)]
#[non_exhaustive]
pub enum TraceContextError {
    /// The `traceparent` string did not have four dash-separated
    /// fields (version, trace-id, parent-id, flags).
    #[error("traceparent must have four dash-separated fields, got `{0}`")]
    BadShape(String),
    /// The `version` field was not the supported `00` value.
    #[error("traceparent version `{0}` is not supported (only `00` is implemented)")]
    UnsupportedVersion(String),
    /// The `trace-id` field was not 32 lowercase hex characters.
    #[error("traceparent trace-id must be 32 lowercase hex chars, got `{0}`")]
    BadTraceId(String),
    /// The `parent-id` (a.k.a. span-id) field was not 16 lowercase
    /// hex characters.
    #[error("traceparent parent-id must be 16 lowercase hex chars, got `{0}`")]
    BadParentId(String),
    /// The `flags` field was not two lowercase hex characters.
    #[error("traceparent flags must be 2 lowercase hex chars, got `{0}`")]
    BadFlags(String),
    /// A `traceparent` with a trace-id or parent-id of all zeros is
    /// invalid per the W3C spec.
    #[error("traceparent carries the all-zero trace-id or parent-id, which is invalid")]
    AllZeroId,
}

/// Sampling and other flag bits from the W3C trace-flags byte.
///
/// Today the only standardized bit is `sampled` (bit 0). The rest is
/// reserved; we hold the raw byte so a future flag can be lifted
/// without breaking the wire shape.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TraceFlags(pub u8);

impl TraceFlags {
    /// Sampling decision: when `true`, downstream is expected to
    /// record the trace.
    pub fn sampled(self) -> bool {
        self.0 & 0x01 != 0
    }
}

/// Parsed W3C `traceparent`.
///
/// Stored as plain fields so callers can reformat or rebase the
/// span-id without re-parsing.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TraceParent {
    /// 16-byte trace identifier, hex-encoded (32 lowercase chars).
    pub trace_id: String,
    /// 8-byte parent span identifier, hex-encoded (16 lowercase
    /// chars). This is the span the new request is a child of.
    pub parent_id: String,
    /// Sampling + reserved bits.
    pub flags: TraceFlags,
}

impl TraceParent {
    /// Parse a W3C `traceparent` string per RFC 9462.
    pub fn parse(s: &str) -> Result<Self, TraceContextError> {
        let parts: Vec<&str> = s.split('-').collect();
        if parts.len() != 4 {
            return Err(TraceContextError::BadShape(s.to_string()));
        }
        let (version, trace_id, parent_id, flags) = (parts[0], parts[1], parts[2], parts[3]);
        if version != "00" {
            return Err(TraceContextError::UnsupportedVersion(version.to_string()));
        }
        if !is_hex_lower(trace_id, 32) {
            return Err(TraceContextError::BadTraceId(trace_id.to_string()));
        }
        if !is_hex_lower(parent_id, 16) {
            return Err(TraceContextError::BadParentId(parent_id.to_string()));
        }
        if !is_hex_lower(flags, 2) {
            return Err(TraceContextError::BadFlags(flags.to_string()));
        }
        if trace_id.chars().all(|c| c == '0') || parent_id.chars().all(|c| c == '0') {
            return Err(TraceContextError::AllZeroId);
        }
        let flags_byte = u8::from_str_radix(flags, 16).map_err(|_| {
            // `is_hex_lower` already validated the chars, so this is
            // unreachable in practice; map to a typed error anyway.
            TraceContextError::BadFlags(flags.to_string())
        })?;
        Ok(Self {
            trace_id: trace_id.to_string(),
            parent_id: parent_id.to_string(),
            flags: TraceFlags(flags_byte),
        })
    }

    /// Format the typed value back into the wire form.
    pub fn to_wire(&self) -> String {
        format!(
            "00-{}-{}-{:02x}",
            self.trace_id, self.parent_id, self.flags.0
        )
    }
}

/// Opaque vendor list from the W3C `tracestate` header.
///
/// Stored as the raw string so we are lossless on round-trip: per
/// the spec the payload is opaque and vendors are free to use any
/// printable ASCII inside their list-member values.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TraceState(pub String);

impl TraceState {
    /// Borrow the wire form.
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Inject the trace context into a request `_meta` block per the
/// W3C convention.
///
/// Always sets `traceparent`. Sets `tracestate` only when the
/// provided value is non-empty (an empty vendor list omits the
/// header by spec). Mutates `meta` in place so the caller can build
/// the full `_meta` map with other mcptest fields and call this last.
pub fn inject_into_meta(
    meta: &mut serde_json::Map<String, Value>,
    parent: &TraceParent,
    state: Option<&TraceState>,
) {
    meta.insert("traceparent".to_string(), Value::String(parent.to_wire()));
    if let Some(s) = state {
        if !s.0.is_empty() {
            meta.insert("tracestate".to_string(), Value::String(s.0.clone()));
        }
    }
}

/// Extract a trace context from a `_meta` block.
///
/// Returns `Ok(None)` when the `_meta` block does not carry a
/// `traceparent` (a non-propagating request is a normal case). A
/// `traceparent` that fails the strict parse returns the typed
/// error so the caller can surface it.
pub fn extract_from_meta(
    meta: &serde_json::Map<String, Value>,
) -> Result<Option<(TraceParent, Option<TraceState>)>, TraceContextError> {
    let Some(parent_value) = meta.get("traceparent").and_then(Value::as_str) else {
        return Ok(None);
    };
    let parent = TraceParent::parse(parent_value)?;
    let state = meta
        .get("tracestate")
        .and_then(Value::as_str)
        .filter(|s| !s.is_empty())
        .map(|s| TraceState(s.to_string()));
    Ok(Some((parent, state)))
}

/// True when `s` is exactly `n` lowercase hex characters.
fn is_hex_lower(s: &str, n: usize) -> bool {
    s.len() == n
        && s.chars()
            .all(|c| c.is_ascii_digit() || ('a'..='f').contains(&c))
}

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

    const SAMPLE_PARENT: &str = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";

    #[test]
    fn parses_a_well_formed_traceparent() {
        let tp = TraceParent::parse(SAMPLE_PARENT).expect("parses");
        assert_eq!(tp.trace_id, "0af7651916cd43dd8448eb211c80319c");
        assert_eq!(tp.parent_id, "b7ad6b7169203331");
        assert_eq!(tp.flags.0, 0x01);
        assert!(tp.flags.sampled());
    }

    #[test]
    fn round_trip_through_wire_is_byte_stable() {
        let tp = TraceParent::parse(SAMPLE_PARENT).unwrap();
        assert_eq!(tp.to_wire(), SAMPLE_PARENT);
    }

    #[test]
    fn rejects_wrong_field_count() {
        let err = TraceParent::parse("00-aaaa-bbbb").unwrap_err();
        assert!(matches!(err, TraceContextError::BadShape(_)));
    }

    #[test]
    fn rejects_unsupported_version() {
        let err = TraceParent::parse("ff-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")
            .unwrap_err();
        match err {
            TraceContextError::UnsupportedVersion(v) => assert_eq!(v, "ff"),
            other => panic!("expected UnsupportedVersion, got {other:?}"),
        }
    }

    #[test]
    fn rejects_uppercase_hex_in_trace_id() {
        let err = TraceParent::parse("00-0AF7651916CD43DD8448EB211C80319C-b7ad6b7169203331-01")
            .unwrap_err();
        assert!(matches!(err, TraceContextError::BadTraceId(_)));
    }

    #[test]
    fn rejects_short_parent_id() {
        let err =
            TraceParent::parse("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169-01").unwrap_err();
        assert!(matches!(err, TraceContextError::BadParentId(_)));
    }

    #[test]
    fn rejects_all_zero_trace_id() {
        let err = TraceParent::parse("00-00000000000000000000000000000000-b7ad6b7169203331-01")
            .unwrap_err();
        assert!(matches!(err, TraceContextError::AllZeroId));
    }

    #[test]
    fn rejects_all_zero_parent_id() {
        let err = TraceParent::parse("00-0af7651916cd43dd8448eb211c80319c-0000000000000000-01")
            .unwrap_err();
        assert!(matches!(err, TraceContextError::AllZeroId));
    }

    #[test]
    fn sampled_flag_round_trips() {
        let unsampled =
            TraceParent::parse("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00").unwrap();
        assert!(!unsampled.flags.sampled());
        assert_eq!(unsampled.to_wire().chars().last(), Some('0'));
    }

    #[test]
    fn inject_writes_both_headers_when_state_is_present() {
        let parent = TraceParent::parse(SAMPLE_PARENT).unwrap();
        let state = TraceState("vendor1=value".to_string());
        let mut meta = serde_json::Map::new();
        inject_into_meta(&mut meta, &parent, Some(&state));
        assert_eq!(meta["traceparent"], json!(SAMPLE_PARENT));
        assert_eq!(meta["tracestate"], json!("vendor1=value"));
    }

    #[test]
    fn inject_omits_tracestate_when_empty() {
        let parent = TraceParent::parse(SAMPLE_PARENT).unwrap();
        let mut meta = serde_json::Map::new();
        inject_into_meta(&mut meta, &parent, Some(&TraceState(String::new())));
        assert!(meta.contains_key("traceparent"));
        assert!(!meta.contains_key("tracestate"));
    }

    #[test]
    fn extract_returns_none_when_traceparent_is_absent() {
        let meta = serde_json::Map::new();
        let result = extract_from_meta(&meta).unwrap();
        assert!(result.is_none());
    }

    #[test]
    fn extract_returns_typed_value_when_present() {
        let mut meta = serde_json::Map::new();
        meta.insert("traceparent".to_string(), json!(SAMPLE_PARENT));
        meta.insert("tracestate".to_string(), json!("vendor1=value"));
        let (parent, state) = extract_from_meta(&meta).unwrap().expect("present");
        assert_eq!(parent.to_wire(), SAMPLE_PARENT);
        assert_eq!(state.unwrap().as_str(), "vendor1=value");
    }

    #[test]
    fn extract_returns_typed_error_on_malformed_traceparent() {
        let mut meta = serde_json::Map::new();
        meta.insert("traceparent".to_string(), json!("totally invalid"));
        let err = extract_from_meta(&meta).unwrap_err();
        assert!(matches!(err, TraceContextError::BadShape(_)));
    }

    #[test]
    fn extract_drops_empty_tracestate() {
        let mut meta = serde_json::Map::new();
        meta.insert("traceparent".to_string(), json!(SAMPLE_PARENT));
        meta.insert("tracestate".to_string(), json!(""));
        let (_, state) = extract_from_meta(&meta).unwrap().expect("parent present");
        assert!(state.is_none());
    }
}
