dexcost 0.9.0

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use rust_decimal::Decimal;
use sha2::{Digest, Sha256};
// parking_lot::RwLock is non-poisoning and runtime-agnostic; lets the
// synchronous pricing API (`get_cost_sync`, `set_custom_pricing_sync`)
// be invoked from inside a Tokio runtime without panicking. The lock
// is only ever held for brief read/write critical sections — no awaits
// happen while holding it — so blocking the executor for the lock
// acquisition is acceptable (Sprint 1 Theme B / B4).
use parking_lot::RwLock;

use crate::core::models::{CostConfidence, PricingSource};

/// Embedded cost map JSON, loaded at compile time.
const COST_MAP_JSON: &str = include_str!("cost_map.json");

/// Result of a cost lookup.
#[derive(Debug, Clone)]
pub struct CostResult {
    pub cost_usd: Decimal,
    pub cost_confidence: CostConfidence,
    pub pricing_source: PricingSource,
    pub pricing_version: String,
}

/// Per-token pricing for a model.
#[derive(Debug, Clone)]
struct ModelPricing {
    input_cost_per_token: Decimal,
    output_cost_per_token: Decimal,
    cache_read_cost: Decimal,
    has_cache_read: bool,
    /// Anthropic-specific rate for tokens *written* to the prompt cache.
    cache_creation_cost: Decimal,
    has_cache_creation: bool,
    /// Anthropic reports cache buckets separately; OpenAI includes cache reads
    /// inside input tokens.
    cache_tokens_are_disjoint: bool,
}

/// Custom per-1k-token pricing set by the user.
#[derive(Debug, Clone)]
struct CustomPricing {
    input_per_1k: Decimal,
    output_per_1k: Decimal,
}

/// Inner mutable state guarded by a parking_lot RwLock.
struct Inner {
    models: HashMap<String, ModelPricing>,
    custom: HashMap<String, CustomPricing>,
    pricing_version: String,
    api_key: Option<String>,
}

/// PricingEngine provides LLM cost lookups from bundled pricing data.
/// Supports exact match, provider-prefix stripping, and date-suffix fallback.
/// Supports background refresh from a control-layer endpoint.
#[derive(Clone)]
pub struct PricingEngine {
    inner: Arc<RwLock<Inner>>,
    stop_signal: Arc<AtomicBool>,
}

impl PricingEngine {
    /// Creates a new PricingEngine from the embedded cost_map.json.
    pub fn new() -> Self {
        Self::from_bytes(COST_MAP_JSON.as_bytes())
    }

    /// Creates a PricingEngine from raw JSON bytes.
    fn from_bytes(data: &[u8]) -> Self {
        let (models, pricing_version) = Self::parse_cost_map(data);
        PricingEngine {
            inner: Arc::new(RwLock::new(Inner {
                models,
                custom: HashMap::new(),
                pricing_version,
                api_key: None,
            })),
            stop_signal: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Parse raw JSON bytes into a model map and version string.
    fn parse_cost_map(data: &[u8]) -> (HashMap<String, ModelPricing>, String) {
        let raw: HashMap<String, serde_json::Value> = match serde_json::from_slice(data) {
            Ok(map) => map,
            Err(e) => {
                eprintln!(
                    "[dexcost] WARNING: failed to parse bundled pricing data: {}",
                    e
                );
                HashMap::new()
            }
        };

        let mut models = HashMap::with_capacity(raw.len());
        for (name, entry) in &raw {
            if let Some(obj) = entry.as_object() {
                let input = Self::decimal_from_value(obj.get("input_cost_per_token"));
                let output = Self::decimal_from_value(obj.get("output_cost_per_token"));
                let (cache_read, has_cache_read) =
                    if let Some(v) = obj.get("cache_read_input_token_cost") {
                        (Self::decimal_from_value(Some(v)), true)
                    } else {
                        (Decimal::ZERO, false)
                    };
                let (cache_creation, has_cache_creation) =
                    if let Some(v) = obj.get("cache_creation_input_token_cost") {
                        (Self::decimal_from_value(Some(v)), true)
                    } else {
                        (Decimal::ZERO, false)
                    };
                let provider = obj
                    .get("litellm_provider")
                    .and_then(serde_json::Value::as_str)
                    .unwrap_or("");

                models.insert(
                    name.clone(),
                    ModelPricing {
                        input_cost_per_token: input,
                        output_cost_per_token: output,
                        cache_read_cost: cache_read,
                        has_cache_read,
                        cache_creation_cost: cache_creation,
                        has_cache_creation,
                        cache_tokens_are_disjoint: Self::uses_disjoint_cache_buckets(
                            name, provider,
                        ),
                    },
                );
            }
        }

        // Compute pricing version as 12-char SHA-256 prefix
        let mut hasher = Sha256::new();
        hasher.update(data);
        let hash = hasher.finalize();
        let pricing_version = hex::encode(&hash[..6]);

        (models, pricing_version)
    }

    /// Parses a Decimal from a JSON value (float or string).
    fn decimal_from_value(v: Option<&serde_json::Value>) -> Decimal {
        match v {
            Some(serde_json::Value::Number(n)) => {
                if let Some(f) = n.as_f64() {
                    Decimal::try_from(f).unwrap_or(Decimal::ZERO)
                } else {
                    Decimal::ZERO
                }
            }
            Some(serde_json::Value::String(s)) => s.parse().unwrap_or(Decimal::ZERO),
            _ => Decimal::ZERO,
        }
    }

    fn uses_disjoint_cache_buckets(model: &str, provider: &str) -> bool {
        let model = model.to_ascii_lowercase();
        let provider = provider.to_ascii_lowercase();
        provider == "anthropic"
            || provider == "vertex_ai-anthropic_models"
            || model.contains("claude")
            || model.contains("anthropic.")
    }

    /// Computes the cost for a model invocation given token counts.
    /// Tries: custom pricing, exact match, provider-prefix stripped, date-suffix fallback.
    ///
    /// `cache_creation_tokens` is the Anthropic-specific count of input tokens
    /// *written* to the prompt cache, charged at the higher
    /// `cache_creation_input_token_cost` rate.
    pub async fn get_cost(
        &self,
        model: &str,
        input_tokens: i64,
        output_tokens: i64,
        cached_tokens: i64,
        cache_creation_tokens: i64,
    ) -> CostResult {
        let inner = self.inner.read();

        // Check custom pricing first
        if let Some(cp) = inner.custom.get(model) {
            return Self::compute_custom_cost(
                cp,
                &inner.pricing_version,
                model,
                input_tokens,
                output_tokens,
                cached_tokens,
                cache_creation_tokens,
            );
        }

        // Try to find model pricing
        if let Some(mp) = Self::find_model_in(&inner.models, model) {
            return Self::compute_cost_from(
                mp,
                &inner.pricing_version,
                input_tokens,
                output_tokens,
                cached_tokens,
                cache_creation_tokens,
            );
        }

        // Unknown model
        CostResult {
            cost_usd: Decimal::ZERO,
            cost_confidence: CostConfidence::Unknown,
            pricing_source: PricingSource::Unknown,
            pricing_version: inner.pricing_version.clone(),
        }
    }

    /// Synchronous cost lookup — acquires a blocking read lock.
    /// Useful when called from non-async context.
    pub fn get_cost_sync(
        &self,
        model: &str,
        input_tokens: i64,
        output_tokens: i64,
        cached_tokens: i64,
        cache_creation_tokens: i64,
    ) -> CostResult {
        let inner = self.inner.read();

        // Check custom pricing first
        if let Some(cp) = inner.custom.get(model) {
            return Self::compute_custom_cost(
                cp,
                &inner.pricing_version,
                model,
                input_tokens,
                output_tokens,
                cached_tokens,
                cache_creation_tokens,
            );
        }

        if let Some(mp) = Self::find_model_in(&inner.models, model) {
            return Self::compute_cost_from(
                mp,
                &inner.pricing_version,
                input_tokens,
                output_tokens,
                cached_tokens,
                cache_creation_tokens,
            );
        }

        CostResult {
            cost_usd: Decimal::ZERO,
            cost_confidence: CostConfidence::Unknown,
            pricing_source: PricingSource::Unknown,
            pricing_version: inner.pricing_version.clone(),
        }
    }

    /// Resolves a model name to its pricing entry using fallback strategies.
    fn find_model_in<'a>(
        models: &'a HashMap<String, ModelPricing>,
        model: &str,
    ) -> Option<&'a ModelPricing> {
        // 1. Exact match
        if let Some(mp) = models.get(model) {
            return Some(mp);
        }

        // 2. Strip provider prefix (e.g. "openai/gpt-4o" -> "gpt-4o")
        if let Some(idx) = model.find('/') {
            let stripped = &model[idx + 1..];
            if let Some(mp) = models.get(stripped) {
                return Some(mp);
            }
        }

        // 3. Longest-prefix walk on `-`: drop trailing components one at a
        //    time and retry. For `"gpt-4o-mini-2024"` this tries
        //    `gpt-4o-mini-2024`, `gpt-4o-mini`, `gpt-4o`, `gpt`.
        //
        //    This handles arbitrary suffixes (date stamps, region tags,
        //    fine-tune labels, etc.), not just the strict `-YYYY-MM-DD`
        //    form. Mirrors the Python SDK `pricing.py` `_resolve_model`
        //    (`pricing.py:248-256`).
        let parts: Vec<&str> = model.split('-').collect();
        for i in (1..parts.len()).rev() {
            let candidate = parts[..i].join("-");
            if let Some(mp) = models.get(&candidate) {
                return Some(mp);
            }
        }

        None
    }

    /// Computes the cost from model pricing and token counts.
    ///
    /// Anthropic reports normal input, cache reads, and cache writes as disjoint
    /// buckets. OpenAI includes cache reads inside `input_tokens`, where the
    /// cached subset must still be clamped and discounted.
    fn compute_cost_from(
        mp: &ModelPricing,
        pricing_version: &str,
        input_tokens: i64,
        output_tokens: i64,
        cached_tokens: i64,
        cache_creation_tokens: i64,
    ) -> CostResult {
        let input_tokens = input_tokens.max(0);
        let output_tokens = output_tokens.max(0);
        let cached_tokens = cached_tokens.max(0);
        let cache_creation_tokens = cache_creation_tokens.max(0);

        let (input_cost, cost_confidence) = if mp.cache_tokens_are_disjoint {
            let mut confidence = CostConfidence::Computed;
            let cache_read_rate = if mp.has_cache_read {
                mp.cache_read_cost
            } else {
                if cached_tokens > 0 {
                    confidence = CostConfidence::Unknown;
                }
                mp.input_cost_per_token
            };
            let cache_creation_rate = if mp.has_cache_creation {
                mp.cache_creation_cost
            } else {
                if cache_creation_tokens > 0 {
                    confidence = CostConfidence::Unknown;
                }
                mp.input_cost_per_token
            };
            (
                mp.input_cost_per_token * Decimal::new(input_tokens, 0)
                    + cache_read_rate * Decimal::new(cached_tokens, 0)
                    + cache_creation_rate * Decimal::new(cache_creation_tokens, 0),
                confidence,
            )
        } else {
            let effective_cached = if mp.has_cache_read {
                cached_tokens.min(input_tokens)
            } else {
                0
            };
            let remaining = input_tokens - effective_cached;
            let effective_creation = if mp.has_cache_creation {
                cache_creation_tokens.min(remaining)
            } else {
                0
            };
            let non_cached = remaining - effective_creation;
            (
                mp.input_cost_per_token * Decimal::new(non_cached, 0)
                    + mp.cache_read_cost * Decimal::new(effective_cached, 0)
                    + mp.cache_creation_cost * Decimal::new(effective_creation, 0),
                CostConfidence::Computed,
            )
        };

        let output_cost = mp.output_cost_per_token * Decimal::new(output_tokens, 0);
        let total = input_cost + output_cost;

        CostResult {
            cost_usd: total,
            cost_confidence,
            pricing_source: PricingSource::Litellm,
            pricing_version: pricing_version.to_string(),
        }
    }

    fn compute_custom_cost(
        pricing: &CustomPricing,
        pricing_version: &str,
        model: &str,
        input_tokens: i64,
        output_tokens: i64,
        cached_tokens: i64,
        cache_creation_tokens: i64,
    ) -> CostResult {
        let thousand = Decimal::new(1000, 0);
        let input_tokens = input_tokens.max(0);
        let output_tokens = output_tokens.max(0);
        let cached_tokens = cached_tokens.max(0);
        let cache_creation_tokens = cache_creation_tokens.max(0);
        let has_unpriced_disjoint_cache = Self::uses_disjoint_cache_buckets(model, "")
            && (cached_tokens > 0 || cache_creation_tokens > 0);
        let billable_input = input_tokens
            + if has_unpriced_disjoint_cache {
                cached_tokens + cache_creation_tokens
            } else {
                0
            };
        let input_cost = pricing.input_per_1k * Decimal::new(billable_input, 0) / thousand;
        let output_cost = pricing.output_per_1k * Decimal::new(output_tokens, 0) / thousand;
        CostResult {
            cost_usd: input_cost + output_cost,
            cost_confidence: if has_unpriced_disjoint_cache {
                CostConfidence::Unknown
            } else {
                CostConfidence::Computed
            },
            pricing_source: PricingSource::Custom,
            pricing_version: pricing_version.to_string(),
        }
    }

    /// Overrides bundled pricing for a model with per-1k-token rates.
    pub async fn set_custom_pricing(
        &self,
        model: &str,
        input_per_1k: Decimal,
        output_per_1k: Decimal,
    ) {
        let mut inner = self.inner.write();
        inner.custom.insert(
            model.to_string(),
            CustomPricing {
                input_per_1k,
                output_per_1k,
            },
        );
    }

    /// Overrides bundled pricing for a model (blocking, non-async).
    pub fn set_custom_pricing_sync(
        &self,
        model: &str,
        input_per_1k: Decimal,
        output_per_1k: Decimal,
    ) {
        let mut inner = self.inner.write();
        inner.custom.insert(
            model.to_string(),
            CustomPricing {
                input_per_1k,
                output_per_1k,
            },
        );
    }

    /// Returns the pricing version (12-char SHA-256 prefix of the cost map data).
    pub fn pricing_version_sync(&self) -> String {
        self.inner.read().pricing_version.clone()
    }

    /// Returns the pricing version asynchronously.
    pub async fn pricing_version(&self) -> String {
        self.inner.read().pricing_version.clone()
    }

    /// Returns the number of models in the pricing data.
    pub fn model_count_sync(&self) -> usize {
        self.inner.read().models.len()
    }

    /// Returns the number of models asynchronously.
    pub async fn model_count(&self) -> usize {
        self.inner.read().models.len()
    }

    /// Configure authentication for control-plane pricing refreshes.
    pub fn set_api_key(&self, api_key: Option<String>) {
        self.inner.write().api_key = api_key;
    }

    // -------------------------------------------------------------------------
    // Background refresh
    // -------------------------------------------------------------------------

    /// Fetch fresh pricing from the control layer. Fail-silent on error.
    ///
    /// Calls the authenticated control-plane pricing endpoint and consumes its
    /// `{ "data": { "pricing_version": "...", "data": { ...models } } }`
    /// response. Invalid or empty catalogs leave the current map untouched.
    pub async fn refresh_from_server(
        &self,
        endpoint: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        Self::refresh_shared(&self.inner, endpoint).await
    }

    /// Start background refresh on a tokio interval.
    ///
    /// Immediately calls `refresh_from_server`, then repeats every `interval`.
    /// Errors from the server are silently ignored (fail-silent).
    /// Only one background task runs at a time; calling this again replaces
    /// the previous stop signal (the old task will also stop on its next tick).
    pub fn start_background_refresh(&self, endpoint: String, interval: Duration) {
        // Reset the stop signal so any previously spawned task (sharing the old
        // Arc) will exit, and our new task starts fresh.
        self.stop_signal.store(false, Ordering::SeqCst);

        let inner = Arc::clone(&self.inner);
        let stop = Arc::clone(&self.stop_signal);

        let Ok(runtime) = tokio::runtime::Handle::try_current() else {
            return;
        };
        runtime.spawn(async move {
            // Build a temporary engine shell that wraps the shared inner state.
            // This avoids cloning all pricing data — we operate directly on the Arc.
            let engine = RefreshWorker {
                inner,
                stop: Arc::clone(&stop),
            };

            // Immediate first refresh
            let _ = engine.refresh(&endpoint).await;

            let mut ticker = tokio::time::interval(interval);
            ticker.tick().await; // consume the first (immediate) tick

            loop {
                ticker.tick().await;
                if stop.load(Ordering::SeqCst) {
                    break;
                }
                let _ = engine.refresh(&endpoint).await;
            }
        });
    }

    /// Stop the background refresh task.
    pub fn stop_background_refresh(&self) {
        self.stop_signal.store(true, Ordering::SeqCst);
    }

    pub(crate) fn background_stop_signal(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.stop_signal)
    }

    async fn refresh_shared(
        inner: &Arc<RwLock<Inner>>,
        endpoint: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let url = format!(
            "{}/v1/api/pricing-data/latest",
            endpoint.trim_end_matches('/')
        );
        let api_key = inner.read().api_key.clone();
        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(10))
            .build()?;
        let mut request = client.get(&url).header("User-Agent", "dexcost-sdk");
        if let Some(key) = api_key.filter(|key| !key.is_empty()) {
            request = request.bearer_auth(key);
        }

        let resp = request.send().await?;
        if !resp.status().is_success() {
            return Err(format!("pricing refresh failed: HTTP {}", resp.status()).into());
        }
        let body = resp.bytes().await?;
        let (new_models, new_version) = Self::parse_control_plane_response(&body)?;

        let mut state = inner.write();
        state.models = new_models;
        state.pricing_version = new_version;
        Ok(())
    }

    fn parse_control_plane_response(
        body: &[u8],
    ) -> Result<(HashMap<String, ModelPricing>, String), Box<dyn std::error::Error + Send + Sync>>
    {
        let invalid =
            |message: &'static str| std::io::Error::new(std::io::ErrorKind::InvalidData, message);
        let outer: serde_json::Value = serde_json::from_slice(body)?;
        let response_data = outer
            .get("data")
            .and_then(serde_json::Value::as_object)
            .ok_or_else(|| invalid("missing 'data' response envelope"))?;
        let mut models = response_data
            .get("data")
            .and_then(serde_json::Value::as_object)
            .cloned()
            .ok_or_else(|| invalid("missing pricing model map"))?;
        models.remove("sample_spec");
        if models.is_empty() {
            return Err(invalid("pricing response contained no models").into());
        }

        let models_bytes = serde_json::to_vec(&models)?;
        let (new_models, fallback_version) = Self::parse_cost_map(&models_bytes);
        if new_models.is_empty() {
            return Err(invalid("pricing response contained no usable models").into());
        }
        let version = response_data
            .get("pricing_version")
            .and_then(serde_json::Value::as_str)
            .filter(|version| !version.is_empty())
            .unwrap_or(&fallback_version)
            .to_string();
        Ok((new_models, version))
    }
}

impl Default for PricingEngine {
    fn default() -> Self {
        Self::new()
    }
}

/// Internal helper used inside the background refresh task.
/// Holds a reference to the shared inner state and the stop flag.
struct RefreshWorker {
    inner: Arc<RwLock<Inner>>,
    stop: Arc<AtomicBool>,
}

impl RefreshWorker {
    async fn refresh(
        &self,
        endpoint: &str,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if self.stop.load(Ordering::SeqCst) {
            return Ok(());
        }

        PricingEngine::refresh_shared(&self.inner, endpoint).await
    }
}

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

    const CONTROL_PLANE_PRICING_RESPONSE: &str =
        include_str!("../../../fixtures/pricing_refresh/control_plane_latest.json");

    #[tokio::test]
    async fn test_pricing_engine_loads() {
        let engine = PricingEngine::new();
        assert!(engine.model_count().await > 0);
        assert!(!engine.pricing_version().await.is_empty());
    }

    #[tokio::test]
    async fn test_known_model_cost() {
        let engine = PricingEngine::new();
        let result = engine.get_cost("gpt-4o", 1000, 500, 0, 0).await;
        if result.cost_confidence == CostConfidence::Computed {
            assert!(result.cost_usd > Decimal::ZERO);
            assert_eq!(result.pricing_source, PricingSource::Litellm);
        }
    }

    #[tokio::test]
    async fn test_unknown_model_cost() {
        let engine = PricingEngine::new();
        let result = engine
            .get_cost("totally-unknown-model-xyz", 1000, 500, 0, 0)
            .await;
        assert_eq!(result.cost_usd, Decimal::ZERO);
        assert_eq!(result.cost_confidence, CostConfidence::Unknown);
        assert_eq!(result.pricing_source, PricingSource::Unknown);
    }

    #[tokio::test]
    async fn test_custom_pricing() {
        let engine = PricingEngine::new();
        let input_per_1k = Decimal::new(1, 3); // 0.001
        let output_per_1k = Decimal::new(2, 3); // 0.002

        engine
            .set_custom_pricing("my-custom-model", input_per_1k, output_per_1k)
            .await;

        let result = engine.get_cost("my-custom-model", 1000, 500, 0, 0).await;
        assert_eq!(result.cost_usd, Decimal::new(2, 3));
        assert_eq!(result.cost_confidence, CostConfidence::Computed);
        assert_eq!(result.pricing_source, PricingSource::Custom);
    }

    #[tokio::test]
    async fn test_custom_pricing_does_not_drop_anthropic_cache_buckets() {
        let engine = PricingEngine::new();
        engine
            .set_custom_pricing("my-claude-model", Decimal::new(1, 3), Decimal::new(2, 3))
            .await;

        let result = engine.get_cost("my-claude-model", 100, 0, 1000, 500).await;
        assert_eq!(result.cost_usd, Decimal::new(16, 4));
        assert_eq!(result.cost_confidence, CostConfidence::Unknown);
        assert_eq!(result.pricing_source, PricingSource::Custom);
    }

    #[tokio::test]
    async fn test_provider_prefix_fallback() {
        let engine = PricingEngine::new();
        let with_prefix = engine.get_cost("openai/gpt-4o", 1000, 500, 0, 0).await;
        let without_prefix = engine.get_cost("gpt-4o", 1000, 500, 0, 0).await;
        assert_eq!(with_prefix.cost_usd, without_prefix.cost_usd);
    }

    #[tokio::test]
    async fn test_zero_tokens() {
        let engine = PricingEngine::new();
        let result = engine.get_cost("gpt-4o", 0, 0, 0, 0).await;
        assert_eq!(result.cost_usd, Decimal::ZERO);
    }

    // Fix 1: a model name with a *non-date* suffix must resolve to its base
    // model's price via the longest-prefix walk, instead of falling through
    // to Unknown (priced at $0). The strict `-YYYY-MM-DD` regex used to miss
    // any suffix that was not a date stamp.
    #[tokio::test]
    async fn test_non_date_suffix_resolves_to_base_price() {
        let engine = PricingEngine::new();

        // Sanity-check: the base model exists in the bundled pricing data.
        let base = engine.get_cost("gpt-4o-mini", 1000, 500, 0, 0).await;
        assert_eq!(
            base.cost_confidence,
            CostConfidence::Computed,
            "base model gpt-4o-mini must be priced"
        );
        assert!(base.cost_usd > Decimal::ZERO);

        // `gpt-4o-mini-2024` is NOT a strict `-YYYY-MM-DD` suffix, so the old
        // regex-only resolver returned Unknown / $0. The longest-prefix walk
        // strips `-2024` and resolves to `gpt-4o-mini`.
        let suffixed = engine.get_cost("gpt-4o-mini-2024", 1000, 500, 0, 0).await;
        assert_eq!(
            suffixed.cost_confidence,
            CostConfidence::Computed,
            "non-date-suffixed model must resolve, not be Unknown"
        );
        assert!(
            suffixed.cost_usd > Decimal::ZERO,
            "non-date-suffixed model must not be priced at $0"
        );
        assert_eq!(
            suffixed.cost_usd, base.cost_usd,
            "suffixed model must inherit the base model's price"
        );

        // The walk also strips multiple trailing components: `gpt-4o-foo-bar`
        // walks down to `gpt-4o`.
        let base_4o = engine.get_cost("gpt-4o", 1000, 500, 0, 0).await;
        let multi = engine.get_cost("gpt-4o-foo-bar", 1000, 500, 0, 0).await;
        assert_eq!(multi.cost_usd, base_4o.cost_usd);
        assert_eq!(multi.cost_confidence, CostConfidence::Computed);
    }

    // -------------------------------------------------------------------------
    // Background refresh tests
    // -------------------------------------------------------------------------

    #[tokio::test]
    async fn test_refresh_from_server_unreachable_fails_silently() {
        let engine = PricingEngine::new();
        // Port 1 is almost certainly not listening — should return an error
        let result = engine.refresh_from_server("http://127.0.0.1:1").await;
        assert!(result.is_err(), "expected error for unreachable endpoint");
    }

    #[tokio::test]
    async fn test_refresh_from_server_bad_json_returns_error() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        // Serve a response with no control-plane data envelope.
        tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                // Drain the HTTP request headers first (required on Windows)
                let mut buf = [0u8; 4096];
                let _ = stream.read(&mut buf).await;
                let body = b"null";
                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                    body.len()
                );
                let _ = stream.write_all(response.as_bytes()).await;
                let _ = stream.write_all(body).await;
            }
        });

        let engine = PricingEngine::new();
        let result = engine
            .refresh_from_server(&format!("http://{}", addr))
            .await;
        // `null` has no pricing envelope — should return an error.
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_refresh_from_server_updates_model_map() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        // Serve the real control-plane response envelope.
        let body = CONTROL_PLANE_PRICING_RESPONSE;
        let body_bytes = body.as_bytes().to_vec();
        let body_len = body_bytes.len();
        let saw_auth = Arc::new(AtomicBool::new(false));
        let saw_auth_server = Arc::clone(&saw_auth);

        tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                // Drain request before responding (required on Windows to avoid ECONNABORTED)
                let mut buf = [0u8; 4096];
                let read = stream.read(&mut buf).await.unwrap_or(0);
                let request = String::from_utf8_lossy(&buf[..read]);
                saw_auth_server.store(
                    request.contains("authorization: Bearer dx_test_refresh")
                        || request.contains("Authorization: Bearer dx_test_refresh"),
                    Ordering::SeqCst,
                );
                let header = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                    body_len
                );
                let _ = stream.write_all(header.as_bytes()).await;
                let _ = stream.write_all(&body_bytes).await;
            }
        });

        let engine = PricingEngine::new();
        engine.set_api_key(Some("dx_test_refresh".to_string()));
        let original_count = engine.model_count().await;

        let result = engine
            .refresh_from_server(&format!("http://{}", addr))
            .await;
        assert!(result.is_ok(), "refresh should succeed: {:?}", result);
        assert!(
            saw_auth.load(Ordering::SeqCst),
            "missing refresh auth header"
        );

        // After refresh the model map should contain only what the server returned
        let new_count = engine.model_count().await;
        assert_eq!(
            new_count, 2,
            "model map should match the shared control-plane fixture"
        );

        // The refreshed model should be queryable
        let cost = engine.get_cost("new-model-v1", 1000, 500, 0, 0).await;
        assert_eq!(cost.cost_confidence, CostConfidence::Computed);
        assert!(cost.cost_usd > Decimal::ZERO);
        // Version should have changed from the bundled one
        assert_eq!(engine.pricing_version().await, "server-v-42");

        assert!(
            original_count != new_count
                || engine.pricing_version().await != PricingEngine::new().pricing_version().await,
            "pricing version should differ from bundled after refresh"
        );
    }

    #[tokio::test]
    async fn test_refresh_from_server_rejects_empty_model_map() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let body = br#"{"data":{"pricing_version":"empty-v1","data":{}}}"#;

        tokio::spawn(async move {
            if let Ok((mut stream, _)) = listener.accept().await {
                let mut buf = [0u8; 4096];
                let _ = stream.read(&mut buf).await;
                let header = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                    body.len()
                );
                let _ = stream.write_all(header.as_bytes()).await;
                let _ = stream.write_all(body).await;
            }
        });

        let engine = PricingEngine::new();
        let original_count = engine.model_count().await;
        let original_version = engine.pricing_version().await;
        let result = engine
            .refresh_from_server(&format!("http://{}", addr))
            .await;

        assert!(result.is_err());
        assert_eq!(engine.model_count().await, original_count);
        assert_eq!(engine.pricing_version().await, original_version);
    }

    #[tokio::test]
    async fn test_start_stop_background_refresh_no_panic() {
        let engine = PricingEngine::new();
        // Start with an unreachable endpoint — refresh errors are swallowed
        engine
            .start_background_refresh("http://127.0.0.1:1".to_string(), Duration::from_millis(50));
        // Give the task a moment to attempt one refresh
        tokio::time::sleep(Duration::from_millis(20)).await;
        // Stop should not panic
        engine.stop_background_refresh();
    }

    #[test]
    fn test_start_background_refresh_without_runtime_is_fail_open() {
        let engine = PricingEngine::new();
        engine
            .start_background_refresh("http://127.0.0.1:1".to_string(), Duration::from_millis(50));
        engine.stop_background_refresh();
    }

    #[tokio::test]
    async fn test_stop_before_start_no_panic() {
        let engine = PricingEngine::new();
        // Calling stop before start should be a no-op
        engine.stop_background_refresh();
    }

    #[tokio::test]
    async fn test_background_refresh_updates_engine() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt};
        use tokio::net::TcpListener;

        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let body = r#"{"data":{"pricing_version":"background-v1","data":{"bg-refresh-model":{"input_cost_per_token":0.005,"output_cost_per_token":0.010}}}}"#;
        let body_bytes = body.as_bytes().to_vec();
        let body_len = body_bytes.len();

        // Serve multiple requests, each time draining the request first
        tokio::spawn(async move {
            loop {
                if let Ok((mut stream, _)) = listener.accept().await {
                    let body_bytes = body_bytes.clone();
                    tokio::spawn(async move {
                        let mut buf = [0u8; 4096];
                        let _ = stream.read(&mut buf).await;
                        let header = format!(
                            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                            body_len
                        );
                        let _ = stream.write_all(header.as_bytes()).await;
                        let _ = stream.write_all(&body_bytes).await;
                    });
                }
            }
        });

        let engine = Arc::new(PricingEngine::new());
        let engine2 = Arc::clone(&engine);

        engine.start_background_refresh(format!("http://{}", addr), Duration::from_millis(50));

        // Wait for at least one refresh cycle
        tokio::time::sleep(Duration::from_millis(300)).await;

        engine2.stop_background_refresh();

        // After background refresh, bg-refresh-model should exist
        let cost = engine2.get_cost("bg-refresh-model", 1000, 500, 0, 0).await;
        assert_eq!(cost.cost_confidence, CostConfidence::Computed);
        assert!(cost.cost_usd > Decimal::ZERO);
    }
}