mcp-usage-export 0.5.0

use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

use mcp_usage_export::{
    AggregatedUsage, BatchExporter, BillingPipeline, MeterEventDeadLetterReason,
    MeterEventExporter, MeterEventOutcome, MeterEventProvider, MeterEventProviderError,
    MeterEventProviderFuture, UsageEvent, UsageRecorder,
};

type ProviderResult = Result<Vec<MeterEventOutcome>, MeterEventProviderError>;

#[derive(Default)]
struct FakeProvider {
    attempts: Mutex<Vec<Vec<AggregatedUsage>>>,
    responses: Mutex<VecDeque<ProviderResult>>,
}

impl FakeProvider {
    fn with_responses(responses: impl IntoIterator<Item = ProviderResult>) -> Self {
        Self {
            attempts: Mutex::new(Vec::new()),
            responses: Mutex::new(responses.into_iter().collect()),
        }
    }

    fn attempts(&self) -> Vec<Vec<AggregatedUsage>> {
        self.attempts
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }
}

impl MeterEventProvider for FakeProvider {
    fn submit<'a>(&'a self, batch: &'a [AggregatedUsage]) -> MeterEventProviderFuture<'a> {
        Box::pin(async move {
            self.attempts
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .push(batch.to_vec());
            self.responses
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .pop_front()
                .unwrap_or_else(|| Ok(vec![MeterEventOutcome::Accepted; batch.len()]))
        })
    }
}

fn aggregate(identifier: &str) -> AggregatedUsage {
    AggregatedUsage {
        identifier: identifier.to_owned(),
        customer_id: format!("customer-{identifier}"),
        meter: format!("meter-{identifier}"),
        units: 42,
        timestamp: 1_800_000_000,
    }
}

#[tokio::test]
async fn accepted_batch_completes_without_retry_progress() {
    let provider = FakeProvider::with_responses([Ok(vec![
        MeterEventOutcome::Accepted,
        MeterEventOutcome::Accepted,
    ])]);
    let exporter = MeterEventExporter::new(provider);
    let batch = [aggregate("one"), aggregate("two")];

    exporter.export(&batch).await.unwrap();

    assert!(!exporter.retry_in_progress());
    assert_eq!(exporter.provider().attempts(), vec![batch.to_vec()]);
}

#[tokio::test]
async fn mixed_outcomes_skip_completed_events_on_retry() {
    let provider = FakeProvider::with_responses([
        Ok(vec![
            MeterEventOutcome::Accepted,
            MeterEventOutcome::PermanentRejection { code: "invalid" },
            MeterEventOutcome::RetryableFailure { code: "busy" },
        ]),
        Ok(vec![MeterEventOutcome::Accepted]),
    ]);
    let exporter = MeterEventExporter::new(provider);
    let batch = [
        aggregate("accepted"),
        aggregate("rejected"),
        aggregate("retryable"),
    ];

    assert!(exporter.export(&batch).await.is_err());
    assert!(exporter.retry_in_progress());
    exporter.export(&batch).await.unwrap();

    let attempts = exporter.provider().attempts();
    assert_eq!(attempts[0], batch);
    assert_eq!(attempts[1], vec![batch[2].clone()]);
    let dead_letters = exporter.take_dead_letters();
    assert_eq!(dead_letters.len(), 1);
    assert_eq!(dead_letters[0].aggregate, batch[1]);
    assert_eq!(
        dead_letters[0].reason,
        MeterEventDeadLetterReason::SynchronousProviderRejection { code: "invalid" }
    );
}

#[tokio::test]
async fn billing_pipeline_retains_original_batch_while_provider_gets_only_unresolved_events() {
    let provider = FakeProvider::with_responses([
        Ok(vec![
            MeterEventOutcome::Accepted,
            MeterEventOutcome::RetryableFailure { code: "busy" },
        ]),
        Ok(vec![MeterEventOutcome::Accepted]),
    ]);
    let pipeline = BillingPipeline::new(MeterEventExporter::new(provider));
    for customer in ["a", "b"] {
        pipeline
            .record(UsageEvent {
                tenant_id: customer.to_owned(),
                customer_id: customer.to_owned(),
                meter: "units".to_owned(),
                units: 1,
                idempotency_key: None,
                timestamp: 1_800_000_000,
            })
            .unwrap();
    }

    assert!(pipeline.flush().await.is_err());
    assert_eq!(pipeline.pending_buckets(), 2);
    assert_eq!(pipeline.flush().await.unwrap(), 2);
    assert_eq!(pipeline.pending_buckets(), 0);

    let attempts = pipeline.exporter().provider().attempts();
    assert_eq!(attempts[0].len(), 2);
    assert_eq!(attempts[1], vec![attempts[0][1].clone()]);
    assert_eq!(attempts[0][1].identifier, attempts[1][0].identifier);
}

#[tokio::test]
async fn batch_failure_retries_every_event_with_stable_identifiers() {
    let provider = FakeProvider::with_responses([
        Err(MeterEventProviderError::new("unavailable")),
        Ok(vec![
            MeterEventOutcome::Accepted,
            MeterEventOutcome::Accepted,
        ]),
    ]);
    let exporter = MeterEventExporter::new(provider);
    let batch = [aggregate("one"), aggregate("two")];

    assert!(exporter.export(&batch).await.is_err());
    exporter.export(&batch).await.unwrap();

    assert_eq!(exporter.provider().attempts(), vec![batch.to_vec(); 2]);
}

#[tokio::test]
async fn malformed_outcome_count_applies_nothing() {
    let provider = FakeProvider::with_responses([
        Ok(vec![MeterEventOutcome::PermanentRejection {
            code: "invalid",
        }]),
        Ok(vec![
            MeterEventOutcome::Accepted,
            MeterEventOutcome::Accepted,
        ]),
    ]);
    let exporter = MeterEventExporter::new(provider);
    let batch = [aggregate("one"), aggregate("two")];

    assert!(exporter.export(&batch).await.is_err());
    assert_eq!(exporter.dead_letter_count(), 0);
    exporter.export(&batch).await.unwrap();

    assert_eq!(exporter.provider().attempts(), vec![batch.to_vec(); 2]);
}

#[tokio::test]
async fn different_batch_is_refused_until_progress_is_abandoned() {
    let provider = FakeProvider::with_responses([
        Ok(vec![
            MeterEventOutcome::Accepted,
            MeterEventOutcome::RetryableFailure { code: "busy" },
        ]),
        Ok(vec![MeterEventOutcome::Accepted]),
    ]);
    let exporter = MeterEventExporter::new(provider);
    let first = [aggregate("complete"), aggregate("unresolved")];
    let different = [aggregate("different")];

    assert!(exporter.export(&first).await.is_err());
    let error = exporter.export(&different).await.unwrap_err();
    assert!(format!("{error}").contains("abandon_retry_progress"));
    assert_eq!(exporter.provider().attempts().len(), 1);
    assert_eq!(exporter.abandon_retry_progress(), vec!["unresolved"]);
    assert!(!exporter.retry_in_progress());
    exporter.export(&different).await.unwrap();
}

#[tokio::test]
async fn full_batch_identity_is_required_for_retry() {
    let provider = FakeProvider::with_responses([Ok(vec![MeterEventOutcome::RetryableFailure {
        code: "busy",
    }])]);
    let exporter = MeterEventExporter::new(provider);
    let original = [aggregate("same-id")];
    let mut changed = original.clone();
    changed[0].units += 1;

    assert!(exporter.export(&original).await.is_err());
    assert!(exporter.export(&changed).await.is_err());
    assert_eq!(exporter.provider().attempts().len(), 1);
}

#[tokio::test]
async fn dead_letters_are_deduplicated_bounded_and_drained() {
    let provider = FakeProvider::with_responses([Ok(vec![
        MeterEventOutcome::PermanentRejection { code: "invalid" },
        MeterEventOutcome::PermanentRejection { code: "invalid" },
        MeterEventOutcome::PermanentRejection { code: "invalid" },
    ])]);
    let exporter = MeterEventExporter::new(provider).with_dead_letter_capacity(1);
    let duplicate = aggregate("duplicate");
    let newest = aggregate("newest");

    exporter
        .export(&[duplicate.clone(), duplicate, newest.clone()])
        .await
        .unwrap();

    assert_eq!(exporter.dead_letter_count(), 1);
    assert_eq!(exporter.dropped_dead_letters(), 1);
    assert_eq!(exporter.take_dead_letters()[0].aggregate, newest);
    assert!(exporter.take_dead_letters().is_empty());
}

#[tokio::test]
async fn zero_capacity_counts_sync_and_async_discards() {
    let provider =
        FakeProvider::with_responses([Ok(vec![MeterEventOutcome::PermanentRejection {
            code: "invalid",
        }])]);
    let exporter = MeterEventExporter::new(provider).with_dead_letter_capacity(0);
    let sync_rejected = aggregate("sync");
    let async_rejected = aggregate("async");

    exporter.export(&[sync_rejected]).await.unwrap();
    assert!(!exporter.quarantine_async_rejection(&async_rejected, "late_invalid"));
    assert_eq!(exporter.dead_letter_count(), 0);
    assert_eq!(exporter.dropped_dead_letters(), 2);
}

#[test]
fn asynchronous_rejections_are_bounded_and_deduplicated() {
    let exporter = MeterEventExporter::new(FakeProvider::default()).with_dead_letter_capacity(1);
    let first = aggregate("first");
    let second = aggregate("second");

    assert!(exporter.quarantine_async_rejection(&first, "late_invalid"));
    assert!(!exporter.quarantine_async_rejection(&first, "late_invalid"));
    assert!(exporter.quarantine_async_rejection(&second, "late_invalid"));
    assert_eq!(exporter.dropped_dead_letters(), 1);
    let retained = exporter.take_dead_letters();
    assert_eq!(retained[0].aggregate, second);
    assert_eq!(
        retained[0].reason,
        MeterEventDeadLetterReason::AsynchronousProviderRejection {
            code: "late_invalid"
        }
    );
}

#[test]
fn debug_output_redacts_provider_and_aggregate_data() {
    struct SensitiveProvider(&'static str);

    impl MeterEventProvider for SensitiveProvider {
        fn submit<'a>(&'a self, batch: &'a [AggregatedUsage]) -> MeterEventProviderFuture<'a> {
            Box::pin(async move { Ok(vec![MeterEventOutcome::Accepted; batch.len()]) })
        }
    }

    let exporter = MeterEventExporter::new(SensitiveProvider("provider-secret"));
    let private = aggregate("private-identifier");
    assert!(exporter.quarantine_async_rejection(&private, "late_invalid"));

    let debug = format!("{exporter:?}");
    for private in [
        "provider-secret",
        "private-identifier",
        "customer-private-identifier",
        "meter-private-identifier",
    ] {
        assert!(!debug.contains(private));
    }
    let dead_letter_debug = format!("{:?}", exporter.take_dead_letters()[0]);
    assert!(!dead_letter_debug.contains("private-identifier"));
    assert_eq!(exporter.provider().0, "provider-secret");
}

#[derive(Default)]
struct CancelOnceProvider {
    attempts: AtomicUsize,
}

impl MeterEventProvider for CancelOnceProvider {
    fn submit<'a>(&'a self, batch: &'a [AggregatedUsage]) -> MeterEventProviderFuture<'a> {
        Box::pin(async move {
            if self.attempts.fetch_add(1, Ordering::SeqCst) == 0 {
                std::future::pending::<()>().await;
            }
            Ok(vec![MeterEventOutcome::Accepted; batch.len()])
        })
    }
}

#[tokio::test]
async fn cancellation_leaves_all_events_retryable() {
    let exporter = Arc::new(MeterEventExporter::new(CancelOnceProvider::default()));
    let batch = vec![aggregate("cancelled")];
    let task = tokio::spawn({
        let exporter = Arc::clone(&exporter);
        let batch = batch.clone();
        async move { exporter.export(&batch).await }
    });
    while exporter.provider().attempts.load(Ordering::SeqCst) == 0 {
        tokio::task::yield_now().await;
    }
    assert!(exporter.export(&batch).await.is_err());
    assert!(exporter.abandon_retry_progress().is_empty());
    assert!(exporter.retry_in_progress());
    task.abort();
    assert!(task.await.unwrap_err().is_cancelled());

    assert!(exporter.retry_in_progress());
    exporter.export(&batch).await.unwrap();
    assert_eq!(exporter.provider().attempts.load(Ordering::SeqCst), 2);
}

#[derive(Default)]
struct PanicOnceProvider {
    panicked: AtomicBool,
    attempts: AtomicUsize,
}

impl MeterEventProvider for PanicOnceProvider {
    fn submit<'a>(&'a self, batch: &'a [AggregatedUsage]) -> MeterEventProviderFuture<'a> {
        Box::pin(async move {
            self.attempts.fetch_add(1, Ordering::SeqCst);
            assert!(
                self.panicked.swap(true, Ordering::SeqCst),
                "intentional provider panic"
            );
            Ok(vec![MeterEventOutcome::Accepted; batch.len()])
        })
    }
}

#[tokio::test]
async fn panic_leaves_all_events_retryable() {
    let exporter = Arc::new(MeterEventExporter::new(PanicOnceProvider::default()));
    let batch = vec![aggregate("panicked")];
    let task = tokio::spawn({
        let exporter = Arc::clone(&exporter);
        let batch = batch.clone();
        async move { exporter.export(&batch).await }
    });
    assert!(task.await.unwrap_err().is_panic());

    assert!(exporter.retry_in_progress());
    exporter.export(&batch).await.unwrap();
    assert_eq!(exporter.provider().attempts.load(Ordering::SeqCst), 2);
}

#[tokio::test]
async fn traffic_during_outage_preserves_batches_through_composite_exporter() {
    let provider = FakeProvider::with_responses([
        Err(MeterEventProviderError::new("unavailable")),
        Err(MeterEventProviderError::new("unavailable")),
    ]);
    let exporter = Arc::new(MeterEventExporter::new(provider));
    let composite = mcp_usage_export::CompositeExporter::new().with_exporter(exporter.clone());
    let pipeline = BillingPipeline::new(composite);
    pipeline
        .record(UsageEvent::now("t", "a", "units", 7, None))
        .unwrap();
    assert!(pipeline.flush().await.is_err());
    pipeline
        .record(UsageEvent::now("t", "a", "units", 11, None))
        .unwrap();
    assert!(pipeline.flush().await.is_err());
    pipeline
        .record(UsageEvent::now("t", "b", "units", 13, None))
        .unwrap();
    assert_eq!(pipeline.flush().await.unwrap(), 1);
    assert_eq!(pipeline.flush().await.unwrap(), 2);
    assert_eq!(pipeline.pending_buckets(), 0);
    let attempts = exporter.provider().attempts();
    assert_eq!(attempts.len(), 4);
    assert_eq!(attempts[0], attempts[1]);
    assert_eq!(attempts[1], attempts[2]);
    assert_eq!(attempts[3].iter().map(|e| e.units).sum::<u64>(), 24);
    assert!(
        attempts[3]
            .iter()
            .all(|e| e.identifier != attempts[0][0].identifier)
    );
}