use std::str::FromStr;

use tonic::metadata::{MetadataMap, MetadataValue};

/// Ordered list of supported distributed trace metadata headers.
const TRACE_METADATA_KEYS: [&str; 8] = [
    "traceparent",
    "tracestate",
    "x-b3-traceid",
    "x-b3-spanid",
    "x-b3-sampled",
    "x-b3-flags",
    "b3",
    "x-request-id",
];

/// Extracted distributed trace context values copied from incoming gRPC metadata.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TraceContext {
    values: Vec<(&'static str, String)>,
}

impl TraceContext {
    /// Extracts supported trace metadata entries from one gRPC metadata map.
    #[must_use]
    pub fn extract(metadata: &MetadataMap) -> Self {
        let mut values = Vec::new();
        for key in TRACE_METADATA_KEYS {
            if let Some(value) = metadata.get(key).and_then(|value| value.to_str().ok()) {
                values.push((key, value.to_string()));
            }
        }
        Self { values }
    }

    /// Injects extracted trace metadata into an outgoing gRPC metadata map.
    pub fn inject(&self, metadata: &mut MetadataMap) {
        for (key, value) in &self.values {
            if let Ok(metadata_value) = MetadataValue::from_str(value.as_str()) {
                metadata.insert(*key, metadata_value);
            }
        }
    }

    /// Returns whether any trace metadata values were captured.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.values.is_empty()
    }

    /// Returns the W3C `traceparent` header value when present.
    #[must_use]
    pub fn traceparent(&self) -> Option<&str> {
        self.values
            .iter()
            .find(|(key, _)| *key == "traceparent")
            .map(|(_, value)| value.as_str())
    }
}

/// Builds a per-request tracing span with propagated trace metadata hints.
#[must_use]
pub fn rpc_span(rpc: &'static str, context: &TraceContext) -> tracing::Span {
    tracing::info_span!(
        "rpc_request",
        rpc,
        traceparent = %context.traceparent().unwrap_or(""),
        has_trace_context = !context.is_empty(),
    )
}

#[cfg(test)]
mod tests {
    use tonic::metadata::{MetadataMap, MetadataValue};

    use super::TraceContext;

    #[test]
    fn extract_and_inject_round_trip_supported_trace_headers() {
        let mut metadata = MetadataMap::new();
        metadata.insert(
            "traceparent",
            MetadataValue::from_static("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"),
        );
        metadata.insert(
            "x-b3-traceid",
            MetadataValue::from_static("4bf92f3577b34da6"),
        );

        let context = TraceContext::extract(&metadata);
        assert!(!context.is_empty());
        assert_eq!(
            context.traceparent(),
            Some("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
        );

        let mut outgoing = MetadataMap::new();
        context.inject(&mut outgoing);
        assert!(outgoing.get("traceparent").is_some());
        assert!(outgoing.get("x-b3-traceid").is_some());
    }
}