use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use anyhow::Result;
use chrono::Utc;
use tokio::sync::Mutex;
use tracing::{debug, warn};

use crate::db::SqliteDb;

/// Tracks LLM and search spending against a configurable daily budget.
///
/// Costs are stored in micro-dollars (1 USD = 1_000_000 units) to avoid
/// floating-point drift while keeping the public API in `f64` dollars.
///
/// Supports:
/// - Configurable per-tier costs (instead of hardcoded values)
/// - SQLite persistence so spend survives restarts
/// - Automatic UTC date rollover detection
pub struct CostOptimizer {
    /// Daily budget in micro-dollars.
    daily_budget_micros: u64,
    /// Running total of today's spend in micro-dollars.
    daily_spend_micros: AtomicU64,
    /// Configurable costs per tier: [tier1, tier2, tier3].
    tier_costs: [f64; 3],
    /// Optional DB for persistence.
    db: Option<Arc<SqliteDb>>,
    /// The current UTC date string (YYYY-MM-DD) for rollover detection.
    current_date: Mutex<String>,
}

const MICROS_PER_DOLLAR: f64 = 1_000_000.0;

impl CostOptimizer {
    /// Create a new optimizer with a daily budget expressed in dollars.
    /// Backward-compatible constructor (no DB, default tier costs).
    pub fn new(daily_budget: f64) -> Self {
        Self::with_config(daily_budget, [0.0, 0.001, 0.01], None)
    }

    /// Create an optimizer with configurable tier costs and optional DB persistence.
    pub fn with_config(
        daily_budget: f64,
        tier_costs: [f64; 3],
        db: Option<Arc<SqliteDb>>,
    ) -> Self {
        let budget_micros = (daily_budget * MICROS_PER_DOLLAR) as u64;
        debug!(daily_budget, budget_micros, "CostOptimizer initialized");
        Self {
            daily_budget_micros: budget_micros,
            daily_spend_micros: AtomicU64::new(0),
            tier_costs,
            db,
            current_date: Mutex::new(Utc::now().format("%Y-%m-%d").to_string()),
        }
    }

    /// Load today's spend from the database. Call after construction.
    pub async fn load_from_db(&self) -> Result<()> {
        if let Some(ref db) = self.db {
            let today = Utc::now().format("%Y-%m-%d").to_string();
            let spend = db.get_daily_spend(&today).await?;
            self.daily_spend_micros.store(spend, Ordering::Relaxed);
            *self.current_date.lock().await = today;
            debug!(spend_micros = spend, "Loaded daily spend from DB");
        }
        Ok(())
    }

    /// Check if the UTC date has rolled over. If so, reset the counter.
    /// Called by the orchestrator before each search to handle day boundaries.
    pub async fn check_date_rollover(&self) {
        let today = Utc::now().format("%Y-%m-%d").to_string();
        let mut current = self.current_date.lock().await;
        if *current != today {
            let old_spend = self.daily_spend_micros.swap(0, Ordering::Relaxed);
            debug!(
                old_date = %*current,
                new_date = %today,
                old_spend_micros = old_spend,
                "Date rollover: daily spend reset"
            );
            *current = today;
        }
    }

    /// Returns `true` if spending `additional_cost` (in dollars) would stay
    /// within the daily budget.
    pub fn can_afford(&self, additional_cost: f64) -> bool {
        let additional_micros = (additional_cost * MICROS_PER_DOLLAR) as u64;
        let current = self.daily_spend_micros.load(Ordering::Relaxed);
        let allowed = current + additional_micros <= self.daily_budget_micros;
        if !allowed {
            warn!(
                current_spend = current as f64 / MICROS_PER_DOLLAR,
                additional_cost,
                budget = self.daily_budget_micros as f64 / MICROS_PER_DOLLAR,
                "Budget check FAILED"
            );
        }
        allowed
    }

    /// Record a cost (in dollars) that has already been incurred.
    /// Persists to SQLite asynchronously if a DB is configured.
    pub fn record_cost(&self, cost: f64) {
        if !cost.is_finite() || cost < 0.0 {
            tracing::warn!(cost, "Ignoring invalid cost value");
            return;
        }
        let micros = (cost * MICROS_PER_DOLLAR) as u64;
        let previous = self.daily_spend_micros.fetch_add(micros, Ordering::Relaxed);
        let new_total = (previous + micros) as f64 / MICROS_PER_DOLLAR;
        debug!(cost, new_total, "Cost recorded");

        if previous + micros > self.daily_budget_micros {
            warn!(
                daily_spend = new_total,
                budget = self.daily_budget_micros as f64 / MICROS_PER_DOLLAR,
                "Daily budget EXCEEDED"
            );
        }

        // Fire-and-forget persistence
        if let Some(ref db) = self.db {
            let db = Arc::clone(db);
            let today = Utc::now().format("%Y-%m-%d").to_string();
            tokio::spawn(async move {
                if let Err(e) = db.add_daily_spend(&today, micros).await {
                    warn!(error = %e, "Failed to persist daily spend");
                }
            });
        }
    }

    /// Return today's accumulated spend in dollars.
    pub fn get_daily_spend(&self) -> f64 {
        self.daily_spend_micros.load(Ordering::Relaxed) as f64 / MICROS_PER_DOLLAR
    }

    /// Return the configured daily budget in dollars.
    pub fn get_daily_budget(&self) -> f64 {
        self.daily_budget_micros as f64 / MICROS_PER_DOLLAR
    }

    /// Return the remaining budget in dollars.
    pub fn get_remaining_budget(&self) -> f64 {
        let spent = self.daily_spend_micros.load(Ordering::Relaxed);
        if spent >= self.daily_budget_micros {
            0.0
        } else {
            (self.daily_budget_micros - spent) as f64 / MICROS_PER_DOLLAR
        }
    }

    /// Reset daily spend back to zero. Should be called at the start of each day.
    pub fn reset_daily_spend(&self) {
        let old = self.daily_spend_micros.swap(0, Ordering::Relaxed);
        debug!(
            previous_spend = old as f64 / MICROS_PER_DOLLAR,
            "Daily spend reset"
        );
    }

    /// Estimate the cost for a given search tier using configured costs.
    pub fn estimate_tier_cost(&self, tier: u8) -> f64 {
        match tier {
            1 => self.tier_costs[0],
            2 => self.tier_costs[1],
            3 => self.tier_costs[2],
            _ => 0.0,
        }
    }

    /// Pick the most cost-effective tier that we can still afford.
    /// Returns `Ok(tier)` or `Err` if even the cheapest tier exceeds budget.
    pub fn recommend_max_tier(&self) -> Result<u8> {
        if self.can_afford(self.estimate_tier_cost(3)) {
            Ok(3)
        } else if self.can_afford(self.estimate_tier_cost(2)) {
            Ok(2)
        } else if self.can_afford(self.estimate_tier_cost(1)) {
            Ok(1)
        } else {
            anyhow::bail!("Daily budget exhausted; cannot afford any search tier")
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db::SqliteDb;
    use tempfile::NamedTempFile;

    #[test]
    fn test_new_optimizer() {
        let opt = CostOptimizer::new(1.0);
        assert!((opt.get_daily_budget() - 1.0).abs() < 1e-6);
        assert!((opt.get_daily_spend() - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_can_afford() {
        let opt = CostOptimizer::new(0.10);
        assert!(opt.can_afford(0.05));
        assert!(opt.can_afford(0.10));
        assert!(!opt.can_afford(0.11));
    }

    #[test]
    fn test_record_cost_and_budget() {
        let opt = CostOptimizer::new(0.10);
        opt.record_cost(0.04);
        assert!((opt.get_daily_spend() - 0.04).abs() < 1e-4);
        assert!(opt.can_afford(0.05));
        opt.record_cost(0.05);
        assert!(!opt.can_afford(0.02));
    }

    #[test]
    fn test_reset() {
        let opt = CostOptimizer::new(1.0);
        opt.record_cost(0.5);
        opt.reset_daily_spend();
        assert!((opt.get_daily_spend() - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_remaining_budget() {
        let opt = CostOptimizer::new(1.0);
        opt.record_cost(0.3);
        let remaining = opt.get_remaining_budget();
        assert!((remaining - 0.7).abs() < 1e-4);
    }

    #[test]
    fn test_recommend_max_tier() {
        let opt = CostOptimizer::new(1.0);
        assert_eq!(opt.recommend_max_tier().unwrap(), 3);

        // Exhaust budget
        opt.record_cost(1.0);
        // Tier 1 is free so should still work
        assert_eq!(opt.recommend_max_tier().unwrap(), 1);
    }

    #[test]
    fn test_configurable_tier_costs() {
        let opt = CostOptimizer::with_config(1.0, [0.05, 0.10, 0.50], None);
        assert!((opt.estimate_tier_cost(1) - 0.05).abs() < f64::EPSILON);
        assert!((opt.estimate_tier_cost(2) - 0.10).abs() < f64::EPSILON);
        assert!((opt.estimate_tier_cost(3) - 0.50).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_date_rollover_resets_spend() {
        let opt = CostOptimizer::new(1.0);
        opt.record_cost(0.5);
        assert!((opt.get_daily_spend() - 0.5).abs() < 1e-4);

        // Simulate date rollover by changing the stored date
        {
            let mut date = opt.current_date.lock().await;
            *date = "1999-01-01".to_string();
        }
        opt.check_date_rollover().await;
        assert!((opt.get_daily_spend() - 0.0).abs() < 1e-6);
    }

    #[tokio::test]
    async fn test_db_persistence_roundtrip() {
        let tmp = NamedTempFile::new().unwrap();
        let db = Arc::new(SqliteDb::new(tmp.path()).unwrap());

        // Create optimizer with DB, record some cost
        let opt1 = CostOptimizer::with_config(1.0, [0.0, 0.001, 0.01], Some(Arc::clone(&db)));
        opt1.record_cost(0.25);
        // Give the fire-and-forget tokio::spawn a moment
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Create a new optimizer and load from DB
        let opt2 = CostOptimizer::with_config(1.0, [0.0, 0.001, 0.01], Some(Arc::clone(&db)));
        opt2.load_from_db().await.unwrap();
        assert!((opt2.get_daily_spend() - 0.25).abs() < 1e-4);
    }

    #[test]
    fn test_backward_compat_no_db() {
        // Default constructor should work fine without DB
        let opt = CostOptimizer::new(0.5);
        opt.record_cost(0.1);
        assert!((opt.get_daily_spend() - 0.1).abs() < 1e-4);
        assert!((opt.estimate_tier_cost(1) - 0.0).abs() < f64::EPSILON);
        assert!((opt.estimate_tier_cost(2) - 0.001).abs() < f64::EPSILON);
        assert!((opt.estimate_tier_cost(3) - 0.01).abs() < f64::EPSILON);
    }
}
