datadog-opentelemetry 0.3.2

// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

use crate::core::configuration::SamplingRuleConfig;
use crate::core::constants::{
    RL_EFFECTIVE_RATE, SAMPLING_AGENT_RATE_TAG_KEY, SAMPLING_DECISION_MAKER_TAG_KEY,
    SAMPLING_KNUTH_RATE_TAG_KEY, SAMPLING_PRIORITY_TAG_KEY, SAMPLING_RULE_RATE_TAG_KEY,
};
use crate::core::sampling::{mechanism, priority, SamplingMechanism, SamplingPriority};

/// Type alias for sampling rules update callback
/// Consolidated callback type used across crates for remote config sampling updates
pub type SamplingRulesCallback = Box<dyn for<'a> Fn(&'a [SamplingRuleConfig]) + Send + Sync>;

use crate::sampling::{AttributeLike, SamplingData, SpanProperties, TraceIdLike, ValueLike};
use std::collections::HashMap;

use super::agent_service_sampler::{AgentRates, ServicesSampler};
// Import the attr constants
use super::constants::pattern::NO_RULE;
use super::glob_matcher::GlobMatcher;
use super::rate_limiter::RateLimiter;
use super::rate_sampler::RateSampler;
use super::rules_sampler::RulesSampler;

fn matcher_from_rule(rule: &str) -> Option<GlobMatcher> {
    (rule != NO_RULE).then(|| GlobMatcher::new(rule))
}

/// Represents a sampling rule with criteria for matching spans
#[derive(Clone, Debug)]
pub struct SamplingRule {
    /// The sample rate to apply when this rule matches (0.0-1.0)
    sample_rate: f64,

    /// Where this rule comes from (customer, dynamic, default)
    provenance: String,

    /// Internal rate sampler used when this rule matches
    rate_sampler: RateSampler,

    /// Glob matchers for pattern matching
    name_matcher: Option<GlobMatcher>,
    service_matcher: Option<GlobMatcher>,
    resource_matcher: Option<GlobMatcher>,
    tag_matchers: HashMap<String, GlobMatcher>,
}

impl SamplingRule {
    /// Converts a vector of SamplingRuleConfig into SamplingRule objects
    /// Centralizes the conversion logic
    pub fn from_configs(configs: Vec<SamplingRuleConfig>) -> Vec<Self> {
        configs
            .into_iter()
            .map(|config| {
                Self::new(
                    config.sample_rate,
                    config.service,
                    config.name,
                    config.resource,
                    Some(config.tags),
                    Some(config.provenance),
                )
            })
            .collect()
    }

    /// Creates a new sampling rule
    pub fn new(
        sample_rate: f64,
        service: Option<String>,
        name: Option<String>,
        resource: Option<String>,
        tags: Option<HashMap<String, String>>,
        provenance: Option<String>,
    ) -> Self {
        // Create glob matchers for the patterns
        let name_matcher = name.as_deref().and_then(matcher_from_rule);
        let service_matcher = service.as_deref().and_then(matcher_from_rule);
        let resource_matcher = resource.as_deref().and_then(matcher_from_rule);

        // Create matchers for tag values
        let tag_map = tags.clone().unwrap_or_default();
        let mut tag_matchers = HashMap::with_capacity(tag_map.len());
        for (key, value) in &tag_map {
            if let Some(matcher) = matcher_from_rule(value) {
                tag_matchers.insert(key.clone(), matcher);
            }
        }

        SamplingRule {
            sample_rate,
            provenance: provenance.unwrap_or_else(|| "default".to_string()),
            rate_sampler: RateSampler::new(sample_rate),
            name_matcher,
            service_matcher,
            resource_matcher,
            tag_matchers,
        }
    }

    /// Checks if this rule matches the given span's attributes and name
    /// The name is derived from the attributes and span kind
    fn matches(&self, span: &impl SpanProperties) -> bool {
        // Get the operation name from the span
        let name = span.operation_name();

        // Check name using glob matcher if specified
        if let Some(ref matcher) = self.name_matcher {
            if !matcher.matches(name.as_ref()) {
                return false;
            }
        }

        // Check service if specified using glob matcher
        if let Some(ref matcher) = self.service_matcher {
            // Get service from the span
            let service = span.service();

            // Match against the service
            if !matcher.matches(&service) {
                return false;
            }
        }

        // Get the resource string for matching
        let resource_str = span.resource();

        // Check resource if specified using glob matcher
        if let Some(ref matcher) = self.resource_matcher {
            // Use the resource from the span
            if !matcher.matches(resource_str.as_ref()) {
                return false;
            }
        }

        // Check all tags using glob matchers
        for (key, matcher) in &self.tag_matchers {
            let rule_tag_key_str = key.as_str();

            // Special handling for rules defined with "http.status_code" or
            // "http.response.status_code"
            if rule_tag_key_str == "http.status_code"
                || rule_tag_key_str
                    == opentelemetry_semantic_conventions::trace::HTTP_RESPONSE_STATUS_CODE
            {
                match self.match_http_status_code_rule(matcher, span) {
                    Some(true) => continue,             // Status code matched
                    Some(false) | None => return false, // Status code didn't match or wasn't found
                }
            } else {
                // Logic for other tags:
                // First, try to match directly with the provided tag key
                let direct_match = span
                    .attributes()
                    .find(|attr| attr.key() == rule_tag_key_str)
                    .and_then(|attr| self.match_attribute_value(attr.value(), matcher));

                if direct_match.unwrap_or(false) {
                    continue;
                }

                // If no direct match, try to find the corresponding OpenTelemetry attribute that
                // maps to the Datadog tag key This handles cases where the rule key
                // is a Datadog key (e.g., "http.method") and the attribute is an
                // OTel key (e.g., "http.request.method")
                if rule_tag_key_str.starts_with("http.") {
                    let tag_match = span.attributes().any(|attr| {
                        if let Some(alternate_key) = span.get_alternate_key(attr.key()) {
                            if alternate_key == rule_tag_key_str {
                                return self
                                    .match_attribute_value(attr.value(), matcher)
                                    .unwrap_or(false);
                            }
                        }
                        false
                    });

                    if !tag_match {
                        return false; // Mapped attribute not found or did not match
                    }
                    // If tag_match is true, loop continues to next rule_tag_key.
                } else {
                    // For non-HTTP attributes, if we don't have a direct match, the rule doesn't
                    // match
                    return false;
                }
            }
        }

        true
    }

    /// Helper method to specifically match a rule against an HTTP status code extracted from
    /// attributes. Returns Some(true) if status code found and matches, Some(false) if found
    /// but not matched, None if not found.
    fn match_http_status_code_rule(
        &self,
        matcher: &GlobMatcher,
        span: &impl SpanProperties,
    ) -> Option<bool> {
        span.status_code().and_then(|status_code| {
            let status_value = opentelemetry::Value::I64(i64::from(status_code));
            self.match_attribute_value(&status_value, matcher)
        })
    }

    // Helper method to match attribute values considering different value types
    fn match_attribute_value(&self, value: &impl ValueLike, matcher: &GlobMatcher) -> Option<bool> {
        // Floating point values are handled with special rules
        if let Some(float_val) = value.extract_float() {
            // Check if the float has a non-zero decimal part
            let has_decimal = float_val != (float_val as i64) as f64;

            // For non-integer floats, only match if it's a wildcard pattern
            if has_decimal {
                // All '*' pattern returns true, any other pattern returns false
                return Some(matcher.pattern().chars().all(|c| c == '*'));
            }

            // For integer floats, convert to string for matching
            return Some(matcher.matches(&float_val.to_string()));
        }

        // For non-float values, use normal matching
        value
            .extract_string()
            .map(|string_value| matcher.matches(&string_value))
    }

    /// Samples a trace ID using this rule's sample rate
    pub fn sample(&self, trace_id: &impl TraceIdLike) -> bool {
        // Delegate to the internal rate sampler's new sample method
        self.rate_sampler.sample(trace_id)
    }
}

/// Represents a priority for sampling rules
#[allow(dead_code)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum RuleProvenance {
    Customer = 0,
    Dynamic = 1,
    Default = 2,
}

impl From<&str> for RuleProvenance {
    fn from(s: &str) -> Self {
        match s {
            "customer" => RuleProvenance::Customer,
            "dynamic" => RuleProvenance::Dynamic,
            _ => RuleProvenance::Default,
        }
    }
}

/// A composite sampler that applies rules in order of precedence
#[derive(Clone, Debug)]
pub struct DatadogSampler {
    /// Sampling rules to apply, in order of precedence
    rules: RulesSampler,

    /// Service-based samplers provided by the Agent
    service_samplers: ServicesSampler,

    /// Rate limiter for limiting the number of spans per second
    rate_limiter: RateLimiter,
}

impl DatadogSampler {
    /// Creates a new DatadogSampler with the given rules
    pub fn new(rules: Vec<SamplingRule>, rate_limit: i32) -> Self {
        // Create rate limiter with default value of 100 if not provided
        let limiter = RateLimiter::new(rate_limit, None);

        DatadogSampler {
            rules: RulesSampler::new(rules),
            service_samplers: ServicesSampler::default(),
            rate_limiter: limiter,
        }
    }

    // used for tests
    #[allow(dead_code)]
    pub(crate) fn update_service_rates(&self, rates: impl IntoIterator<Item = (String, f64)>) {
        self.service_samplers.update_rates(rates);
    }

    pub(crate) fn on_agent_response(&self) -> Box<dyn for<'a> Fn(&'a str) + Send + Sync> {
        let service_samplers = self.service_samplers.clone();
        Box::new(move |s: &str| {
            let Ok(new_rates) = serde_json::de::from_str::<AgentRates>(s) else {
                return;
            };
            let Some(new_rates) = new_rates.rates_by_service else {
                return;
            };
            service_samplers.update_rates(new_rates.into_iter().map(|(k, v)| (k.to_string(), v)));
        })
    }

    /// Creates a callback for updating sampling rules from remote configuration
    /// # Returns
    /// A boxed function that takes a slice of SamplingRuleConfig and updates the sampling rules
    pub fn on_rules_update(&self) -> SamplingRulesCallback {
        let rules_sampler = self.rules.clone();
        Box::new(move |rule_configs: &[SamplingRuleConfig]| {
            let new_rules = SamplingRule::from_configs(rule_configs.to_vec());

            rules_sampler.update_rules(new_rules);
        })
    }

    /// Computes a key for service-based sampling
    fn service_key(&self, span: &impl SpanProperties) -> String {
        // Get service from span
        let service = span.service().into_owned();
        // Get env from span
        let env = span.env();

        format!("service:{service},env:{env}")
    }

    /// Finds the highest precedence rule that matches the span
    fn find_matching_rule(&self, span: &impl SpanProperties) -> Option<SamplingRule> {
        self.rules.find_matching_rule(|rule| rule.matches(span))
    }

    /// Returns the sampling mechanism used for the decision
    fn get_sampling_mechanism(
        &self,
        rule: Option<&SamplingRule>,
        used_agent_sampler: bool,
    ) -> SamplingMechanism {
        if let Some(rule) = rule {
            match rule.provenance.as_str() {
                // Provenance will not be set for rules until we implement remote configuration
                "customer" => mechanism::REMOTE_USER_TRACE_SAMPLING_RULE,
                "dynamic" => mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE,
                _ => mechanism::LOCAL_USER_TRACE_SAMPLING_RULE,
            }
        } else if used_agent_sampler {
            // If using service-based sampling from the agent
            mechanism::AGENT_RATE_BY_SERVICE
        } else {
            // Should not happen, but just in case
            mechanism::DEFAULT
        }
    }

    /// Sample an incoming span based on the parent context and attributes
    pub fn sample(&self, data: &impl SamplingData) -> DdSamplingResult {
        if let Some(is_parent_sampled) = data.is_parent_sampled() {
            let priority = match is_parent_sampled {
                false => priority::AUTO_REJECT,
                true => priority::AUTO_KEEP,
            };
            // If a parent exists, inherit its sampling decision and trace state
            return DdSamplingResult {
                priority,
                trace_root_info: None,
            };
        }

        // Apply rules-based sampling
        data.with_span_properties(self, |sampler, span| sampler.sample_root(data, span))
    }

    /// Sample the root span of a trace
    fn sample_root(
        &self,
        data: &impl SamplingData,
        span: &impl SpanProperties,
    ) -> DdSamplingResult {
        let mut is_keep = true;
        let mut used_agent_sampler = false;
        let sample_rate;
        let mut rl_effective_rate: Option<f64> = None;
        let trace_id = data.trace_id();

        // Find a matching rule
        let matching_rule = self.find_matching_rule(span);

        // Apply sampling logic
        if let Some(rule) = &matching_rule {
            // Get the sample rate from the rule
            sample_rate = rule.sample_rate;

            // First check if the span should be sampled according to the rule
            if !rule.sample(trace_id) {
                is_keep = false;
            // If the span should be sampled, then apply rate limiting
            } else if !self.rate_limiter.is_allowed() {
                is_keep = false;
                rl_effective_rate = Some(self.rate_limiter.effective_rate());
            }
        } else {
            // Try service-based sampling from Agent
            let service_key = self.service_key(span);
            if let Some(sampler) = self.service_samplers.get(&service_key) {
                // Use the service-based sampler
                used_agent_sampler = true;
                sample_rate = sampler.sample_rate(); // Get rate for reporting

                // Check if the service sampler decides to drop
                if !sampler.sample(trace_id) {
                    is_keep = false;
                }
            } else {
                // Default sample rate, should never happen in practice if agent provides rates
                sample_rate = 1.0;
                // Keep the default decision (RecordAndSample)
            }
        }

        // Determine the sampling mechanism
        let mechanism = self.get_sampling_mechanism(matching_rule.as_ref(), used_agent_sampler);

        DdSamplingResult {
            priority: mechanism.to_priority(is_keep),
            trace_root_info: Some(TraceRootSamplingInfo {
                mechanism,
                rate: sample_rate,
                rl_effective_rate,
            }),
        }
    }
}

/// Formats a sampling rate with up to 6 significant digits, stripping trailing zeros.
///
/// This matches the Go behavior of `strconv.FormatFloat(rate, 'g', 6, 64)`.
///
/// # Examples
/// - `1.0` → `Some("1")`
/// - `0.5` → `Some("0.5")`
/// - `0.7654321` → `Some("0.765432")`
/// - `0.100000` → `Some("0.1")`
/// - `-0.1` → `None`
/// - `1.1` → `None`
fn format_sampling_rate(rate: f64) -> Option<String> {
    if rate.is_nan() || !(0.0..=1.0).contains(&rate) {
        return None;
    }

    if rate == 0.0 {
        return Some("0".to_string());
    }

    let digits = 6_i32;
    let magnitude = rate.abs().log10().floor() as i32;
    let scale = 10f64.powi(digits - 1 - magnitude);
    let rounded = (rate * scale).round() / scale;

    // Determine decimal places needed for 6 significant digits
    let decimal_places = if magnitude >= digits - 1 {
        0
    } else {
        (digits - 1 - magnitude) as usize
    };

    let s = format!("{:.prec$}", rounded, prec = decimal_places);
    // Strip trailing zeros after decimal point
    Some(if s.contains('.') {
        let s = s.trim_end_matches('0');
        let s = s.trim_end_matches('.');
        s.to_string()
    } else {
        s
    })
}

pub struct TraceRootSamplingInfo {
    mechanism: SamplingMechanism,
    rate: f64,
    rl_effective_rate: Option<f64>,
}

impl TraceRootSamplingInfo {
    /// Returns the sampling mechanism used for this trace root
    pub fn mechanism(&self) -> SamplingMechanism {
        self.mechanism
    }

    /// Returns the sample rate used for this trace root
    pub fn rate(&self) -> f64 {
        self.rate
    }

    /// Returns the effective rate limit if rate limiting was applied
    pub fn rl_effective_rate(&self) -> Option<f64> {
        self.rl_effective_rate
    }
}

pub struct DdSamplingResult {
    priority: SamplingPriority,
    trace_root_info: Option<TraceRootSamplingInfo>,
}

impl DdSamplingResult {
    #[inline(always)]
    pub fn get_priority(&self) -> SamplingPriority {
        self.priority
    }

    pub fn get_trace_root_sampling_info(&self) -> &Option<TraceRootSamplingInfo> {
        &self.trace_root_info
    }

    /// Returns Datadog-specific sampling tags to be added as attributes
    ///
    /// # Parameters
    /// * `factory` - The attribute factory to use for creating attributes
    ///
    /// # Returns
    /// An optional vector of attributes to add to the sampling result
    pub fn to_dd_sampling_tags<F>(&self, factory: &F) -> Option<Vec<F::Attribute>>
    where
        F: crate::sampling::AttributeFactory,
    {
        let Some(root_info) = &self.trace_root_info else {
            return None; // No root info, return empty attributes
        };

        let mut result: Vec<F::Attribute>;
        // Add rate limiting tag if applicable
        if let Some(limit) = root_info.rl_effective_rate() {
            result = Vec::with_capacity(4);
            result.push(factory.create_f64(RL_EFFECTIVE_RATE, limit));
        } else {
            result = Vec::with_capacity(3);
        }

        // Add the sampling decision trace tag with the mechanism
        let mechanism = root_info.mechanism();
        result.push(factory.create_string(SAMPLING_DECISION_MAKER_TAG_KEY, mechanism.to_cow()));

        // Add the sample rate tag with the correct key based on the mechanism
        match mechanism {
            mechanism::AGENT_RATE_BY_SERVICE => {
                result.push(factory.create_f64(SAMPLING_AGENT_RATE_TAG_KEY, root_info.rate()));
                if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
                    result.push(factory.create_string(
                        SAMPLING_KNUTH_RATE_TAG_KEY,
                        std::borrow::Cow::Owned(rate_str),
                    ));
                }
            }
            mechanism::REMOTE_USER_TRACE_SAMPLING_RULE
            | mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE
            | mechanism::LOCAL_USER_TRACE_SAMPLING_RULE => {
                result.push(factory.create_f64(SAMPLING_RULE_RATE_TAG_KEY, root_info.rate()));
                if let Some(rate_str) = format_sampling_rate(root_info.rate()) {
                    result.push(factory.create_string(
                        SAMPLING_KNUTH_RATE_TAG_KEY,
                        std::borrow::Cow::Owned(rate_str),
                    ));
                }
            }
            _ => {}
        }

        let priority = self.priority;
        result.push(factory.create_i64(SAMPLING_PRIORITY_TAG_KEY, priority.into_i8() as i64));

        Some(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mappings::get_otel_operation_name_v2;
    use crate::sampling::constants::{
        attr::{ENV_TAG, RESOURCE_TAG},
        pattern,
    };
    use crate::sampling::otel_mappings::{OtelSamplingData, PreSampledSpan};
    use opentelemetry::trace::{SpanKind, TraceId};
    use opentelemetry::{Key, KeyValue, Value};
    use opentelemetry_sdk::Resource as SdkResource;
    use opentelemetry_semantic_conventions::{
        attribute::{
            DB_SYSTEM_NAME, HTTP_REQUEST_METHOD, MESSAGING_OPERATION_TYPE, MESSAGING_SYSTEM,
        },
        resource::SERVICE_NAME,
        trace::{HTTP_RESPONSE_STATUS_CODE, NETWORK_PROTOCOL_NAME},
    };
    use std::sync::{Arc, RwLock};

    fn create_empty_resource() -> opentelemetry_sdk::Resource {
        opentelemetry_sdk::Resource::builder_empty().build()
    }

    // Helper function to create an empty resource wrapped in Arc<RwLock> for DatadogSampler
    fn create_empty_resource_arc() -> Arc<RwLock<opentelemetry_sdk::Resource>> {
        Arc::new(RwLock::new(
            opentelemetry_sdk::Resource::builder_empty().build(),
        ))
    }

    fn create_resource(res: String) -> Arc<RwLock<SdkResource>> {
        let attributes = vec![
            KeyValue::new(SERVICE_NAME, res), // String `res` is Into<Value>
        ];
        let resource: SdkResource = SdkResource::builder_empty()
            .with_attributes(attributes)
            .build();
        Arc::new(RwLock::new(resource))
    }

    // Helper function to create a trace ID
    fn create_trace_id() -> TraceId {
        let bytes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
        TraceId::from_bytes(bytes)
    }

    // Helper function to create attributes for testing
    fn create_attributes(resource: &'static str, env: &'static str) -> Vec<KeyValue> {
        vec![
            KeyValue::new(RESOURCE_TAG, resource),
            KeyValue::new("datadog.env", env),
        ]
    }

    // Helper function to create SamplingData for testing
    fn create_sampling_data<'a>(
        is_parent_sampled: Option<bool>,
        trace_id: &'a TraceId,
        name: &'a str,
        span_kind: SpanKind,
        attributes: &'a [KeyValue],
        resource: &'a RwLock<opentelemetry_sdk::Resource>,
    ) -> OtelSamplingData<'a> {
        OtelSamplingData::new(
            is_parent_sampled,
            trace_id,
            name,
            span_kind,
            attributes,
            resource,
        )
    }

    #[test]
    fn test_sampling_rule_creation() {
        let rule = SamplingRule::new(
            0.5,
            Some("test-service".to_string()),
            Some("test-name".to_string()),
            Some("test-resource".to_string()),
            Some(HashMap::from([(
                "custom-tag".to_string(),
                "tag-value".to_string(),
            )])),
            Some("customer".to_string()),
        );

        assert_eq!(rule.sample_rate, 0.5);
        assert_eq!(rule.service_matcher.unwrap().pattern(), "test-service");
        assert_eq!(rule.name_matcher.unwrap().pattern(), "test-name");
        assert_eq!(
            rule.resource_matcher.unwrap().pattern(),
            "test-resource".to_string()
        );
        assert_eq!(
            rule.tag_matchers.get("custom-tag").unwrap().pattern(),
            "tag-value"
        );
        assert_eq!(rule.provenance, "customer");
    }

    #[test]
    fn test_sampling_rule_with_no_rule() {
        // Create a rule without specifying any criteria
        let rule = SamplingRule::new(
            0.5, None, // No service
            None, // No name
            None, // No resource
            None, // No tags
            None, // Default provenance
        );

        // Verify fields are set to None or empty
        assert_eq!(rule.sample_rate, 0.5);
        assert!(rule.service_matcher.is_none());
        assert!(rule.name_matcher.is_none());
        assert!(rule.resource_matcher.is_none());
        assert!(rule.tag_matchers.is_empty());
        assert_eq!(rule.provenance, "default");

        // Verify no matchers were created
        assert!(rule.service_matcher.is_none());
        assert!(rule.name_matcher.is_none());
        assert!(rule.resource_matcher.is_none());
        assert!(rule.tag_matchers.is_empty());

        // Test that a rule with NO_RULE constants behaves the same as None
        let rule_with_empty_strings = SamplingRule::new(
            0.5,
            Some(pattern::NO_RULE.to_string()), // Empty service string
            Some(pattern::NO_RULE.to_string()), // Empty name string
            Some(pattern::NO_RULE.to_string()), // Empty resource string
            Some(HashMap::from([(
                pattern::NO_RULE.to_string(),
                pattern::NO_RULE.to_string(),
            )])), // Empty tag
            None,
        );

        // Verify that matchers aren't created for NO_RULE values
        assert!(rule_with_empty_strings.service_matcher.is_none());
        assert!(rule_with_empty_strings.name_matcher.is_none());
        assert!(rule_with_empty_strings.resource_matcher.is_none());
        assert!(rule_with_empty_strings.tag_matchers.is_empty());

        // Create a span with some attributes
        let attributes = create_attributes("some-resource", "some-env");

        // Empty resource for testing (unwrapped for the test)
        let empty_resource = create_empty_resource();

        // Both rules should match any span since they have no criteria
        let span = PreSampledSpan::new("", SpanKind::Client, &attributes, &empty_resource);
        assert!(rule.matches(&span));
        assert!(rule_with_empty_strings.matches(&span));
    }

    #[test]
    fn test_sampling_rule_matches() {
        // Create a rule with specific service and name patterns
        let _rule = SamplingRule::new(
            0.5,
            Some("web-*".to_string()),
            Some("http.*".to_string()),
            None,
            Some(HashMap::from([(
                "custom_key".to_string(),
                "custom_value".to_string(),
            )])),
            None,
        );
    }

    #[test]
    fn test_sample_method() {
        // Create two rules with different rates
        let rule_always = SamplingRule::new(1.0, None, None, None, None, None);
        let rule_never = SamplingRule::new(0.0, None, None, None, None, None);

        let trace_id = create_trace_id();

        // Rule with rate 1.0 should always sample
        assert!(rule_always.sample(&trace_id));

        // Rule with rate 0.0 should never sample
        assert!(!rule_never.sample(&trace_id));
    }

    #[test]
    fn test_datadog_sampler_creation() {
        // Create a sampler with default config
        let sampler = DatadogSampler::new(vec![], 100);
        assert!(sampler.rules.is_empty());
        assert!(sampler.service_samplers.is_empty());

        // Create a sampler with rules
        let rule = SamplingRule::new(0.5, None, None, None, None, None);
        let sampler_with_rules = DatadogSampler::new(vec![rule], 200);
        assert_eq!(sampler_with_rules.rules.len(), 1);
    }

    #[test]
    fn test_service_key_generation() {
        // Create resource with test service name
        let test_service_name = "test-service".to_string();
        let resource = create_resource(test_service_name.clone());
        let sampler = DatadogSampler::new(vec![], 100);

        // Test with service and env
        // The 'service' in create_attributes is not used for the service part of the key,
        // but ENV_TAG is still correctly picked up from attributes.
        let attrs = create_attributes("resource", "production");
        let res = &resource.read().unwrap();
        let span = PreSampledSpan::new("test-span", SpanKind::Internal, attrs.as_slice(), res);
        assert_eq!(
            sampler.service_key(&span),
            // Expect the service name from the resource
            format!("service:{test_service_name},env:production")
        );

        // Test with missing env
        // The 'service' in these attributes is also not used for the service part of the key.
        let attrs_no_env = vec![KeyValue::new(RESOURCE_TAG, "resource")];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Internal,
            attrs_no_env.as_slice(),
            res,
        );
        assert_eq!(
            sampler.service_key(&span),
            // Expect the service name from the resource and an empty env
            format!("service:{test_service_name},env:")
        );
    }

    #[test]
    fn test_update_service_rates() {
        let sampler = DatadogSampler::new(vec![], 100);

        // Update with service rates
        let mut rates = HashMap::new();
        rates.insert("service:web,env:prod".to_string(), 0.5);
        rates.insert("service:api,env:prod".to_string(), 0.75);

        sampler.service_samplers.update_rates(rates);

        // Check number of samplers
        assert_eq!(sampler.service_samplers.len(), 2);

        // Verify keys exist
        assert!(sampler
            .service_samplers
            .contains_key("service:web,env:prod"));
        assert!(sampler
            .service_samplers
            .contains_key("service:api,env:prod"));

        // Verify the sampling rates are correctly set
        if let Some(web_sampler) = sampler.service_samplers.get("service:web,env:prod") {
            assert_eq!(web_sampler.sample_rate(), 0.5);
        } else {
            panic!("Web service sampler not found");
        }

        if let Some(api_sampler) = sampler.service_samplers.get("service:api,env:prod") {
            assert_eq!(api_sampler.sample_rate(), 0.75);
        } else {
            panic!("API service sampler not found");
        }
    }

    #[test]
    fn test_find_matching_rule() {
        // Create rules with different priorities and service matchers
        let rule1 = SamplingRule::new(
            0.1,
            Some("service1".to_string()),
            None,
            None,
            None,
            Some("customer".to_string()), // Highest priority
        );

        let rule2 = SamplingRule::new(
            0.2,
            Some("service2".to_string()),
            None,
            None,
            None,
            Some("dynamic".to_string()), // Middle priority
        );

        let rule3 = SamplingRule::new(
            0.3,
            Some("service*".to_string()), // Wildcard service
            None,
            None,
            None,
            Some("default".to_string()), // Lowest priority
        );

        let sampler = DatadogSampler::new(vec![rule1.clone(), rule2.clone(), rule3.clone()], 100);

        // Test with a specific service that should match the first rule (rule1)
        {
            let resource = create_resource("service1".to_string());
            let attrs1 = create_attributes("resource_val_for_attr1", "prod");
            let res = resource.read().unwrap();
            let span = PreSampledSpan::new("test-span", SpanKind::Client, attrs1.as_slice(), &res);
            let matching_rule_for_attrs1 = sampler.find_matching_rule(&span);
            assert!(
                matching_rule_for_attrs1.is_some(),
                "Expected rule1 to match for service1"
            );
            let rule = matching_rule_for_attrs1.unwrap();
            assert_eq!(rule.sample_rate, 0.1, "Expected rule1 sample rate");
            assert_eq!(rule.provenance, "customer", "Expected rule1 provenance");
        }

        // Test with a specific service that should match the second rule (rule2)
        {
            let resource = create_resource("service2".to_string());
            let attrs2 = create_attributes("resource_val_for_attr2", "prod");
            let res = resource.read().unwrap();
            let span = PreSampledSpan::new("test-span", SpanKind::Client, attrs2.as_slice(), &res);
            let matching_rule_for_attrs2 = sampler.find_matching_rule(&span);
            assert!(
                matching_rule_for_attrs2.is_some(),
                "Expected rule2 to match for service2"
            );
            let rule = matching_rule_for_attrs2.unwrap();
            assert_eq!(rule.sample_rate, 0.2, "Expected rule2 sample rate");
            assert_eq!(rule.provenance, "dynamic", "Expected rule2 provenance");
        }

        // Test with a service that matches the wildcard rule (rule3)
        {
            let resource = create_resource("service3".to_string());
            let attrs3 = create_attributes("resource_val_for_attr3", "prod");
            let res = resource.read().unwrap();
            let span = PreSampledSpan::new("test-span", SpanKind::Client, attrs3.as_slice(), &res);
            let matching_rule_for_attrs3 = sampler.find_matching_rule(&span);
            assert!(
                matching_rule_for_attrs3.is_some(),
                "Expected rule3 to match for service3"
            );
            let rule = matching_rule_for_attrs3.unwrap();
            assert_eq!(rule.sample_rate, 0.3, "Expected rule3 sample rate");
            assert_eq!(rule.provenance, "default", "Expected rule3 provenance");
        }

        // Test with a service that doesn't match any rule's service pattern
        {
            let resource = create_resource("other_sampler_service".to_string());
            let attrs4 = create_attributes("resource_val_for_attr4", "prod");
            let res = resource.read().unwrap();
            let span = PreSampledSpan::new("test-span", SpanKind::Client, attrs4.as_slice(), &res);
            let matching_rule_for_attrs4 = sampler.find_matching_rule(&span);
            assert!(
                matching_rule_for_attrs4.is_none(),
                "Expected no rule to match for service 'other_sampler_service'"
            );
        }
    }

    #[test]
    fn test_get_sampling_mechanism() {
        let sampler = DatadogSampler::new(vec![], 100);

        // Create rules with different provenances
        let rule_customer =
            SamplingRule::new(0.1, None, None, None, None, Some("customer".to_string()));
        let rule_dynamic =
            SamplingRule::new(0.2, None, None, None, None, Some("dynamic".to_string()));
        let rule_default =
            SamplingRule::new(0.3, None, None, None, None, Some("default".to_string()));

        // Test with customer rule
        let mechanism1 = sampler.get_sampling_mechanism(Some(&rule_customer), false);
        assert_eq!(mechanism1, mechanism::REMOTE_USER_TRACE_SAMPLING_RULE);

        // Test with dynamic rule
        let mechanism2 = sampler.get_sampling_mechanism(Some(&rule_dynamic), false);
        assert_eq!(mechanism2, mechanism::REMOTE_DYNAMIC_TRACE_SAMPLING_RULE);

        // Test with default rule
        let mechanism3 = sampler.get_sampling_mechanism(Some(&rule_default), false);
        assert_eq!(mechanism3, mechanism::LOCAL_USER_TRACE_SAMPLING_RULE);

        // Test with agent sampler
        let mechanism4 = sampler.get_sampling_mechanism(None, true);
        assert_eq!(mechanism4, mechanism::AGENT_RATE_BY_SERVICE);

        // Test fallback case
        let mechanism5 = sampler.get_sampling_mechanism(None, false);
        assert_eq!(mechanism5, mechanism::DEFAULT);
    }

    #[test]
    fn test_add_dd_sampling_tags() {
        // Test with RecordAndSample decision and LocalUserTraceSamplingRule mechanism
        let sample_rate = 0.5;
        let is_sampled = true;
        let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
        let sampling_result = DdSamplingResult {
            priority: mechanism.to_priority(is_sampled),
            trace_root_info: Some(TraceRootSamplingInfo {
                mechanism,
                rate: 0.5,
                rl_effective_rate: None,
            }),
        };

        let attrs = sampling_result
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .unwrap_or_default();

        // Verify the number of attributes (decision_maker + priority + rule_rate + ksr)
        assert_eq!(attrs.len(), 4);

        // Check individual attributes
        let mut found_decision_maker = false;
        let mut found_priority = false;
        let mut found_rule_rate = false;
        let mut found_ksr = false;

        for attr in &attrs {
            match attr.key.as_str() {
                SAMPLING_DECISION_MAKER_TAG_KEY => {
                    let value_str = match &attr.value {
                        opentelemetry::Value::String(s) => s.to_string(),
                        _ => panic!("Expected string value for decision maker tag"),
                    };
                    assert_eq!(value_str, mechanism.to_cow());
                    found_decision_maker = true;
                }
                SAMPLING_PRIORITY_TAG_KEY => {
                    // For LocalUserTraceSamplingRule with KEEP, it should be USER_KEEP
                    let expected_priority = mechanism.to_priority(true).into_i8() as i64;

                    let value_int = match attr.value {
                        opentelemetry::Value::I64(i) => i,
                        _ => panic!("Expected integer value for priority tag"),
                    };
                    assert_eq!(value_int, expected_priority);
                    found_priority = true;
                }
                SAMPLING_RULE_RATE_TAG_KEY => {
                    let value_float = match attr.value {
                        opentelemetry::Value::F64(f) => f,
                        _ => panic!("Expected float value for rule rate tag"),
                    };
                    assert_eq!(value_float, sample_rate);
                    found_rule_rate = true;
                }
                SAMPLING_KNUTH_RATE_TAG_KEY => {
                    let value_str = match &attr.value {
                        opentelemetry::Value::String(s) => s.to_string(),
                        _ => panic!("Expected string value for ksr tag"),
                    };
                    assert_eq!(value_str, "0.5");
                    found_ksr = true;
                }
                _ => {}
            }
        }

        assert!(found_decision_maker, "Missing decision maker tag");
        assert!(found_priority, "Missing priority tag");
        assert!(found_rule_rate, "Missing rule rate tag");
        assert!(found_ksr, "Missing knuth sampling rate tag");

        // Test with rate limiting
        let rate_limit = 0.5;
        let is_sampled = false;
        let mechanism = mechanism::LOCAL_USER_TRACE_SAMPLING_RULE;
        let sampling_result = DdSamplingResult {
            priority: mechanism.to_priority(is_sampled),
            trace_root_info: Some(TraceRootSamplingInfo {
                mechanism,
                rate: 0.5,
                rl_effective_rate: Some(rate_limit),
            }),
        };
        let attrs_with_limit = sampling_result
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .unwrap_or_default();

        // With rate limiting, there should be one more attribute
        assert_eq!(attrs_with_limit.len(), 5);

        // Check for rate limit attribute
        let mut found_limit = false;
        for attr in &attrs_with_limit {
            if attr.key.as_str() == RL_EFFECTIVE_RATE {
                let value_float = match attr.value {
                    opentelemetry::Value::F64(f) => f,
                    _ => panic!("Expected float value for rate limit tag"),
                };
                assert_eq!(value_float, rate_limit);
                found_limit = true;
                break;
            }
        }

        assert!(found_limit, "Missing rate limit tag");

        // Test with AgentRateByService mechanism to check for SAMPLING_AGENT_RATE_TAG_KEY

        let agent_rate = 0.75;
        let is_sampled = false;
        let mechanism = mechanism::AGENT_RATE_BY_SERVICE;
        let sampling_result = DdSamplingResult {
            priority: mechanism.to_priority(is_sampled),
            trace_root_info: Some(TraceRootSamplingInfo {
                mechanism,
                rate: agent_rate,
                rl_effective_rate: None,
            }),
        };

        let agent_attrs = sampling_result
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .unwrap_or_default();

        // Verify the number of attributes (should be 4: decision_maker + priority +
        // agent_rate + ksr)
        assert_eq!(agent_attrs.len(), 4);

        // Check for agent rate tag and ksr tag
        let mut found_agent_rate = false;
        let mut found_ksr = false;
        for attr in &agent_attrs {
            match attr.key.as_str() {
                SAMPLING_AGENT_RATE_TAG_KEY => {
                    let value_float = match attr.value {
                        opentelemetry::Value::F64(f) => f,
                        _ => panic!("Expected float value for agent rate tag"),
                    };
                    assert_eq!(value_float, agent_rate);
                    found_agent_rate = true;
                }
                SAMPLING_KNUTH_RATE_TAG_KEY => {
                    let value_str = match &attr.value {
                        opentelemetry::Value::String(s) => s.to_string(),
                        _ => panic!("Expected string value for ksr tag"),
                    };
                    assert_eq!(value_str, "0.75");
                    found_ksr = true;
                }
                _ => {}
            }
        }

        assert!(found_agent_rate, "Missing agent rate tag");
        assert!(
            found_ksr,
            "Missing knuth sampling rate tag for agent mechanism"
        );

        // Also check that the SAMPLING_RULE_RATE_TAG_KEY is NOT present for agent mechanism
        for attr in &agent_attrs {
            assert_ne!(
                attr.key.as_str(),
                SAMPLING_RULE_RATE_TAG_KEY,
                "Rule rate tag should not be present for agent mechanism"
            );
        }
    }

    #[test]
    fn test_format_sampling_rate() {
        // Exact values
        assert_eq!(format_sampling_rate(1.0), Some("1".to_string()));
        assert_eq!(format_sampling_rate(0.5), Some("0.5".to_string()));
        assert_eq!(format_sampling_rate(0.1), Some("0.1".to_string()));
        assert_eq!(format_sampling_rate(0.0), Some("0".to_string()));

        // Trailing zeros should be stripped
        assert_eq!(format_sampling_rate(0.100000), Some("0.1".to_string()));
        assert_eq!(format_sampling_rate(0.500000), Some("0.5".to_string()));

        // Truncation to 6 significant digits
        assert_eq!(
            format_sampling_rate(0.7654321),
            Some("0.765432".to_string())
        );
        assert_eq!(
            format_sampling_rate(0.123456789),
            Some("0.123457".to_string())
        );

        // Small values
        assert_eq!(format_sampling_rate(0.001), Some("0.001".to_string()));

        // Boundary values
        assert_eq!(format_sampling_rate(0.75), Some("0.75".to_string()));
        assert_eq!(format_sampling_rate(0.999999), Some("0.999999".to_string()));

        // Invalid rates
        assert_eq!(format_sampling_rate(-0.1), None);
        assert_eq!(format_sampling_rate(1.1), None);
        assert_eq!(format_sampling_rate(f64::NAN), None);
        assert_eq!(format_sampling_rate(f64::INFINITY), None);
        assert_eq!(format_sampling_rate(f64::NEG_INFINITY), None);
    }

    #[test]
    fn test_should_sample_parent_context() {
        let sampler = DatadogSampler::new(vec![], 100);

        // Create empty slices for attributes and links
        let empty_attrs: &[KeyValue] = &[];
        let trace_id = create_trace_id();
        let span_kind = SpanKind::Client;
        let resource = create_empty_resource_arc();

        // Test with sampled parent context
        let data_sampled = create_sampling_data(
            Some(true),
            &trace_id,
            "span",
            span_kind.clone(),
            empty_attrs,
            resource.as_ref(),
        );
        let result_sampled = sampler.sample(&data_sampled);

        // Should inherit the sampling decision from parent
        assert!(result_sampled.get_priority().is_keep());
        assert!(result_sampled
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .is_none());

        // Test with non-sampled parent context
        let data_not_sampled = create_sampling_data(
            Some(false),
            &trace_id,
            "span",
            span_kind,
            empty_attrs,
            resource.as_ref(),
        );
        let result_not_sampled = sampler.sample(&data_not_sampled);

        // Should inherit the sampling decision from parent
        assert!(!result_not_sampled.get_priority().is_keep());
        assert!(result_not_sampled
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .is_none());
    }

    #[test]
    fn test_should_sample_with_rule() {
        // Create a rule that always samples
        let rule = SamplingRule::new(
            1.0,
            Some("test-service".to_string()),
            None,
            None,
            None,
            None,
        );

        let sampler = DatadogSampler::new(vec![rule], 100);

        let trace_id = create_trace_id();
        let span_kind = SpanKind::Client;
        let resource = create_empty_resource_arc();

        // Test with matching attributes
        let attrs = create_attributes("resource", "prod");
        let data = create_sampling_data(
            None,
            &trace_id,
            "span",
            span_kind.clone(),
            attrs.as_slice(),
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should sample and add attributes
        assert!(result.get_priority().is_keep());
        assert!(result
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .is_some());

        // Test with non-matching attributes
        let attrs_no_match = create_attributes("other-resource", "prod");
        let data_no_match = create_sampling_data(
            None,
            &trace_id,
            "span",
            span_kind,
            attrs_no_match.as_slice(),
            resource.as_ref(),
        );
        let result_no_match = sampler.sample(&data_no_match);

        // Should still sample (default behavior when no rules match) and add attributes
        assert!(result_no_match.get_priority().is_keep());
        assert!(result_no_match
            .to_dd_sampling_tags(&crate::sampling::OtelAttributeFactory)
            .is_some());
    }

    #[test]
    fn test_should_sample_with_service_rates() {
        // Initialize sampler
        let sampler = DatadogSampler::new(vec![], 100);

        // Add service rates for different service+env combinations
        let mut rates = HashMap::new();
        rates.insert("service:test-service,env:prod".to_string(), 1.0); // Always sample for test-service in prod
        rates.insert("service:other-service,env:prod".to_string(), 0.0); // Never sample for other-service in prod

        sampler.update_service_rates(rates);

        let trace_id = create_trace_id();
        let span_kind = SpanKind::Client;

        // Test with attributes that should lead to "service:test-service,env:prod" key
        let resource_test_service = create_resource("test-service".to_string());
        let attrs_sample = create_attributes("any_resource_name_matching_env", "prod");
        let data_sample = create_sampling_data(
            None,
            &trace_id,
            "span_for_test_service",
            span_kind.clone(),
            attrs_sample.as_slice(),
            resource_test_service.as_ref(),
        );
        let result_sample = sampler.sample(&data_sample);
        // Expect RecordAndSample because service_key will be "service:test-service,env:prod" ->
        // rate 1.0
        assert!(
            result_sample.get_priority().is_keep(),
            "Span for test-service/prod should be sampled"
        );

        // Test with attributes that should lead to "service:other-service,env:prod" key
        let resource_other_service = create_resource("other-service".to_string());
        let attrs_no_sample = create_attributes("any_resource_name_matching_env", "prod");
        let data_no_sample = create_sampling_data(
            None,
            &trace_id,
            "span_for_other_service",
            span_kind,
            attrs_no_sample.as_slice(),
            resource_other_service.as_ref(),
        );
        let result_no_sample = sampler.sample(&data_no_sample);
        // Expect Drop because service_key will be "service:other-service,env:prod" -> rate 0.0
        assert!(
            !result_no_sample.get_priority().is_keep(),
            "Span for other-service/prod should be dropped"
        );
    }

    #[test]
    fn test_sampling_rule_matches_float_attributes() {
        use opentelemetry::Value;

        // Helper to create attributes with a float value
        fn create_attributes_with_float(tag_key: &'static str, float_value: f64) -> Vec<KeyValue> {
            vec![
                KeyValue::new(RESOURCE_TAG, "resource"),
                KeyValue::new(ENV_TAG, "prod"),
                KeyValue::new(tag_key, Value::F64(float_value)),
            ]
        }

        // Test case 1: Rule with exact value matching integer float
        let rule_integer = SamplingRule::new(
            0.5,
            None,
            None,
            None,
            Some(HashMap::from([("float_tag".to_string(), "42".to_string())])),
            None,
        );

        // Should match integer float
        let integer_float_attrs = create_attributes_with_float("float_tag", 42.0);
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            integer_float_attrs.as_slice(),
            &resource,
        );
        assert!(rule_integer.matches(&span));

        // Test case 2: Rule with wildcard pattern and non-integer float
        let rule_wildcard = SamplingRule::new(
            0.5,
            None,
            None,
            None,
            Some(HashMap::from([("float_tag".to_string(), "*".to_string())])),
            None,
        );

        // Should match non-integer float with wildcard pattern
        let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            decimal_float_attrs.as_slice(),
            &resource,
        );
        assert!(rule_wildcard.matches(&span));

        // Test case 3: Rule with specific pattern and non-integer float
        // With our simplified logic, non-integer floats will never match non-wildcard patterns
        let rule_specific = SamplingRule::new(
            0.5,
            None,
            None,
            None,
            Some(HashMap::from([(
                "float_tag".to_string(),
                "42.5".to_string(),
            )])),
            None,
        );

        // Should NOT match the exact decimal value because non-integer floats only match wildcards
        let decimal_float_attrs = create_attributes_with_float("float_tag", 42.5);
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            decimal_float_attrs.as_slice(),
            &resource,
        );
        assert!(!rule_specific.matches(&span));
        // Test case 4: Pattern with partial wildcard '*' for suffix
        let rule_prefix = SamplingRule::new(
            0.5,
            None,
            None,
            None,
            Some(HashMap::from([(
                "float_tag".to_string(),
                "42.*".to_string(),
            )])),
            None,
        );

        // Should NOT match decimal values as we don't do partial pattern matching for non-integer
        // floats
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            decimal_float_attrs.as_slice(),
            &resource,
        );
        assert!(!rule_prefix.matches(&span));
    }

    #[test]
    fn test_otel_to_datadog_attribute_mapping() {
        // Test with a rule that matches against a Datadog attribute name
        let rule = SamplingRule::new(
            1.0,
            None,
            None,
            None,
            Some(HashMap::from([(
                "http.response.status_code".to_string(),
                "5*".to_string(),
            )])),
            None,
        );

        // Create attributes with OpenTelemetry naming convention
        let otel_attrs = vec![KeyValue::new("http.response.status_code", 500)];

        // The rule should match because both use the same OpenTelemetry attribute name
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            otel_attrs.as_slice(),
            &resource,
        );
        assert!(rule.matches(&span));

        // Attributes that don't match the value pattern shouldn't match
        let non_matching_attrs = vec![KeyValue::new("http.response.status_code", 200)];
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            non_matching_attrs.as_slice(),
            &resource,
        );
        assert!(!rule.matches(&span));

        // Attributes that have no mapping to the rule tag shouldn't match
        let unrelated_attrs = vec![KeyValue::new("unrelated.attribute", "value")];
        let resource = create_empty_resource();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            unrelated_attrs.as_slice(),
            &resource,
        );
        assert!(!rule.matches(&span));
    }

    #[test]
    fn test_multiple_otel_attribute_mappings() {
        // Test with a rule that has multiple tag criteria
        let mut tags = HashMap::new();
        tags.insert("http.status_code".to_string(), "5*".to_string());
        tags.insert("http.method".to_string(), "POST".to_string());
        tags.insert("http.url".to_string(), "*api*".to_string());

        let rule = SamplingRule::new(1.0, None, None, None, Some(tags), None);

        // Create attributes with mixed OpenTelemetry and Datadog naming
        let mixed_attrs = vec![
            // OTel attribute that maps to http.status_code
            KeyValue::new("http.response.status_code", 503),
            // OTel attribute that maps to http.method
            KeyValue::new("http.request.method", "POST"),
            // OTel attribute that maps to http.url
            KeyValue::new("url.full", "https://example.com/api/v1/resource"),
        ];

        // The rule should match because all three criteria are satisfied through mapping
        let resource = create_empty_resource();
        let span = PreSampledSpan::new("test-span", SpanKind::Client, &mixed_attrs, &resource);
        assert!(rule.matches(&span));

        // If any criteria is not met, the rule shouldn't match
        let missing_method = vec![
            KeyValue::new("http.response.status_code", 503),
            // Missing http.method/http.request.method
            KeyValue::new("url.full", "https://example.com/api/v1/resource"),
        ];

        let resource = create_empty_resource();
        let span = PreSampledSpan::new("test-span", SpanKind::Client, &missing_method, &resource);
        assert!(!rule.matches(&span));

        // Wrong value should also not match
        let wrong_method = vec![
            KeyValue::new("http.response.status_code", 503),
            KeyValue::new("http.request.method", "GET"), // Not POST
            KeyValue::new("url.full", "https://example.com/api/v1/resource"),
        ];

        let resource = create_empty_resource();
        let span = PreSampledSpan::new("test-span", SpanKind::Client, &wrong_method, &resource);
        assert!(!rule.matches(&span));
    }

    #[test]
    fn test_direct_and_mapped_mixed_attributes() {
        // Constants for key names to improve readability and ensure consistency
        let dd_status_key_str = HTTP_RESPONSE_STATUS_CODE;
        let otel_response_status_key_str = HTTP_RESPONSE_STATUS_CODE;
        let custom_tag_key = "custom.tag";
        let custom_tag_value = "value";

        let empty_resource = create_empty_resource();
        let span_kind_client = SpanKind::Client;

        // Test with both direct matches and mapped attributes
        let mut tags_rule1 = HashMap::new();
        tags_rule1.insert(dd_status_key_str.to_string(), "5*".to_string());
        tags_rule1.insert(custom_tag_key.to_string(), custom_tag_value.to_string());

        let rule1 = SamplingRule::new(1.0, None, None, None, Some(tags_rule1), None);

        // Case 1: OTel attribute that maps to http.status_code (503 matches "5*") + Direct
        // custom.tag match
        let mixed_attrs_match = vec![
            KeyValue::new(otel_response_status_key_str, 503),
            KeyValue::new(custom_tag_key, custom_tag_value),
        ];
        let span = PreSampledSpan::new(
            "test-span",
            span_kind_client,
            &mixed_attrs_match,
            &empty_resource,
        );
        assert!(rule1.matches(&span), "Rule with dd_status_key (5*) and custom.tag should match span with otel_response_status_key (503) and custom.tag");

        // Case 2: Datadog convention for status code (503 matches "5*") + Direct custom.tag match
        let dd_attrs_match = vec![
            KeyValue::new(dd_status_key_str, 503),
            KeyValue::new(custom_tag_key, custom_tag_value),
        ];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            &dd_attrs_match,
            &empty_resource,
        );
        assert!(rule1.matches(&span), "Rule with dd_status_key (5*) and custom.tag should match span with dd_status_key (503) and custom.tag");

        // Case 3: Missing the custom tag should fail (status code would match)
        let missing_custom_tag_attrs = vec![KeyValue::new(otel_response_status_key_str, 503)];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            &missing_custom_tag_attrs,
            &empty_resource,
        );
        assert!(
            !rule1.matches(&span),
            "Rule with dd_status_key (5*) and custom.tag should NOT match span missing custom.tag"
        );

        // Case 4: OTel status code 200 (does NOT match "5*") + custom.tag present
        let non_matching_otel_status_attrs = vec![
            KeyValue::new(otel_response_status_key_str, 200),
            KeyValue::new(custom_tag_key, custom_tag_value),
        ];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            &non_matching_otel_status_attrs,
            &empty_resource,
        );
        assert!(!rule1.matches(&span), "Rule with dd_status_key (5*) and custom.tag should NOT match span with non-matching otel_response_status_key (200)");

        // Case 5: No recognizable status code + custom.tag present
        let no_status_code_attrs = vec![
            KeyValue::new("another.tag", "irrelevant"),
            KeyValue::new(custom_tag_key, custom_tag_value),
        ];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            &no_status_code_attrs,
            &empty_resource,
        );
        assert!(!rule1.matches(&span), "Rule with dd_status_key (5*) and custom.tag should NOT match span with no status code attribute");

        // Case 6: Rule uses OTel key http.response.status_code directly, span has matching OTel
        // key.
        let mut tags_rule2 = HashMap::new();
        tags_rule2.insert(otel_response_status_key_str.to_string(), "200".to_string());
        tags_rule2.insert(custom_tag_key.to_string(), custom_tag_value.to_string());
        let rule2 = SamplingRule::new(1.0, None, None, None, Some(tags_rule2), None);

        let otel_key_rule_match_attrs = vec![
            KeyValue::new(otel_response_status_key_str, 200),
            KeyValue::new(custom_tag_key, custom_tag_value),
        ];
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            &otel_key_rule_match_attrs,
            &empty_resource,
        );
        assert!(rule2.matches(&span), "Rule with otel_response_status_key (200) and custom.tag should match span with otel_response_status_key (200) and custom.tag");
    }

    #[test]
    fn test_operation_name_integration() {
        // Create rules that match different operation name patterns
        let http_rule = SamplingRule::new(
            1.0,                                // 100% sample rate
            None,                               // no service matcher
            Some("http.*.request".to_string()), // matches both client and server HTTP requests
            None,                               // no resource matcher
            None,                               // no tag matchers
            Some("default".to_string()),        // rule name - default provenance
        );

        let db_rule = SamplingRule::new(
            1.0,                                  // 100% sample rate
            None,                                 // no service matcher
            Some("postgresql.query".to_string()), // matches database queries
            None,                                 // no resource matcher
            None,                                 // no tag matchers
            Some("default".to_string()),          // rule name - default provenance
        );

        let messaging_rule = SamplingRule::new(
            1.0,                               // 100% sample rate
            None,                              // no service matcher
            Some("kafka.process".to_string()), // matches Kafka messaging operations
            None,                              // no resource matcher
            None,                              // no tag matchers
            Some("default".to_string()),       // rule name - default provenance
        );

        // Create a sampler with these rules
        let sampler = DatadogSampler::new(vec![http_rule, db_rule, messaging_rule], 100);

        // Create a trace ID for testing
        let trace_id = create_trace_id();
        let resource = create_empty_resource_arc();

        // Test cases for different span kinds and attributes

        // 1. HTTP client request
        let http_client_attrs = vec![KeyValue::new(
            Key::from_static_str(HTTP_REQUEST_METHOD),
            Value::String("GET".into()),
        )];

        let empty_resource: SdkResource = create_empty_resource();
        // Print the operation name that will be generated
        let span = PreSampledSpan::new("", SpanKind::Client, &http_client_attrs, &empty_resource);
        let http_client_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            http_client_op_name, "http.client.request",
            "HTTP client operation name should be correct"
        );

        let span_kind_client = SpanKind::Client;
        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_client.clone(),
            &http_client_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should be sampled due to matching the http_rule
        assert!(result.get_priority().is_keep());

        // 2. HTTP server request
        let http_server_attrs = vec![KeyValue::new(
            Key::from_static_str(HTTP_REQUEST_METHOD),
            Value::String("POST".into()),
        )];

        // Print the operation name that will be generated
        let span = PreSampledSpan::new("", SpanKind::Server, &http_server_attrs, &empty_resource);
        let http_server_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            http_server_op_name, "http.server.request",
            "HTTP server operation name should be correct"
        );

        let span_kind_server = SpanKind::Server;
        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_server.clone(),
            &http_server_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should be sampled due to matching the http_rule
        assert!(result.get_priority().is_keep());

        // 3. Database query
        let db_attrs = vec![KeyValue::new(
            Key::from_static_str(DB_SYSTEM_NAME),
            Value::String("postgresql".into()),
        )];

        // Print the operation name that will be generated
        let span = PreSampledSpan::new("", SpanKind::Client, &db_attrs, &empty_resource);
        let db_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            db_op_name, "postgresql.query",
            "Database operation name should be correct"
        );

        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_client, // DB queries use client span kind
            &db_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should be sampled due to matching the db_rule
        assert!(result.get_priority().is_keep());

        // 4. Messaging operation
        let messaging_attrs = vec![
            KeyValue::new(
                Key::from_static_str(MESSAGING_SYSTEM),
                Value::String("kafka".into()),
            ),
            KeyValue::new(
                Key::from_static_str(MESSAGING_OPERATION_TYPE),
                Value::String("process".into()),
            ),
        ];

        // Print the operation name that will be generated
        let span = PreSampledSpan::new("", SpanKind::Consumer, &messaging_attrs, &empty_resource);
        let messaging_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            messaging_op_name, "kafka.process",
            "Messaging operation name should be correct"
        );

        let span_kind_consumer = SpanKind::Consumer;
        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_consumer, // Messaging uses consumer span kind
            &messaging_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should be sampled due to matching the messaging_rule
        assert!(result.get_priority().is_keep());

        // 5. Generic internal span (should not match any rules)
        let internal_attrs = vec![KeyValue::new("custom.tag", "value")];

        // Print the operation name that will be generated
        let span = PreSampledSpan::new("", SpanKind::Internal, &internal_attrs, &empty_resource);
        let internal_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            internal_op_name, "internal",
            "Internal operation name should be the span kind"
        );

        let span_kind_internal = SpanKind::Internal;
        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_internal,
            &internal_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should still be sampled (default behavior when no rules match)
        assert!(result.get_priority().is_keep());

        // 6. Server with protocol but no HTTP method
        let server_protocol_attrs = vec![KeyValue::new(
            Key::from_static_str(NETWORK_PROTOCOL_NAME),
            Value::String("http".into()),
        )];

        // Print the operation name that will be generated
        let span = PreSampledSpan::new(
            "",
            SpanKind::Server,
            &server_protocol_attrs,
            &empty_resource,
        );
        let server_protocol_op_name = get_otel_operation_name_v2(&span);
        assert_eq!(
            server_protocol_op_name, "http.server.request",
            "Server with protocol operation name should use protocol"
        );

        let data = create_sampling_data(
            None,
            &trace_id,
            "test-span",
            span_kind_server,
            &server_protocol_attrs,
            resource.as_ref(),
        );
        let result = sampler.sample(&data);

        // Should not match our http rule since operation name would be "http.server.request"
        // But should still be sampled (default behavior)
        assert!(result.get_priority().is_keep());
    }

    #[test]
    fn test_on_rules_update_callback() {
        // Create a sampler with initial rules
        let initial_rule = SamplingRule::new(
            0.1,
            Some("initial-service".to_string()),
            None,
            None,
            None,
            Some("default".to_string()),
        );

        // Create a resource with a service name that will match our test rule
        let test_resource = Arc::new(RwLock::new(
            opentelemetry_sdk::Resource::builder_empty()
                .with_attributes(vec![KeyValue::new(SERVICE_NAME, "web-frontend")])
                .build(),
        ));

        let sampler = DatadogSampler::new(vec![initial_rule], 100);

        // Verify initial state
        assert_eq!(sampler.rules.len(), 1);

        // Get the callback
        let callback = sampler.on_rules_update();

        // Create new rules directly as SamplingRuleConfig objects
        let new_rules = vec![
            SamplingRuleConfig {
                sample_rate: 0.5,
                service: Some("web-*".to_string()),
                name: Some("http.*".to_string()),
                resource: None,
                tags: std::collections::HashMap::new(),
                provenance: "customer".to_string(),
            },
            SamplingRuleConfig {
                sample_rate: 0.2,
                service: Some("api-*".to_string()),
                name: None,
                resource: Some("/api/*".to_string()),
                tags: [("env".to_string(), "prod".to_string())].into(),
                provenance: "dynamic".to_string(),
            },
        ];

        // Apply the update
        callback(&new_rules);

        // Verify the rules were updated
        assert_eq!(sampler.rules.len(), 2);

        // Test that the new rules work by finding a matching rule
        // Create attributes that will generate an operation name matching "http.*"
        let attrs = vec![
            KeyValue::new(HTTP_REQUEST_METHOD, "GET"), /* This will make operation name
                                                        * "http.client.request" */
        ];
        let resource_guard = test_resource.read().unwrap();
        let span = PreSampledSpan::new(
            "test-span",
            SpanKind::Client,
            attrs.as_slice(),
            &resource_guard,
        );

        let matching_rule = sampler.find_matching_rule(&span);
        assert!(matching_rule.is_some(), "Expected to find a matching rule for service 'web-frontend' and name 'http.client.request'");
        let rule = matching_rule.unwrap();
        assert_eq!(rule.sample_rate, 0.5);
        assert_eq!(rule.provenance, "customer");

        // Test with empty rules array
        callback(&[]);
        assert_eq!(sampler.rules.len(), 0); // Should now have no rules
    }
}