use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Model identifier (e.g., "anthropic/claude-sonnet-4-20250514")
    pub model: String,
    /// Input/prompt tokens
    pub input_tokens: u64,
    /// Output/completion tokens
    pub output_tokens: u64,
    /// Cached input tokens (Anthropic `cache_read_input_tokens`, OpenAI
    /// `prompt_tokens_details.cached_tokens`). Subset of `input_tokens`
    /// when reported by the provider; the rate sheet's
    /// `cached_input_per_mtok` applies to these.
    #[serde(default, skip_serializing_if = "is_zero_u64")]
    pub cached_input_tokens: u64,
    /// Total tokens (input + output, ignoring the cached subset).
    pub total_tokens: u64,
    /// Calculated cost in USD
    pub cost_usd: f64,
    #[serde(default = "default_true", skip_serializing_if = "is_true_bool")]
    pub pricing_available: bool,
    /// Timestamp of the request
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

fn is_zero_u64(v: &u64) -> bool {
    *v == 0
}

fn default_true() -> bool {
    true
}

fn is_true_bool(v: &bool) -> bool {
    *v
}

impl TokenUsage {
    fn sanitize_price(value: f64) -> f64 {
        if value.is_finite() && value > 0.0 {
            value
        } else {
            0.0
        }
    }

    pub fn billable_input_tokens(&self) -> u64 {
        self.input_tokens.saturating_sub(self.cached_input_tokens)
    }

    pub fn new(
        model: impl Into<String>,
        input_tokens: u64,
        output_tokens: u64,
        cached_input_tokens: u64,
        input_price_per_million: f64,
        output_price_per_million: f64,
        cached_input_price_per_million: f64,
    ) -> Self {
        Self::new_with_cache(
            model,
            input_tokens,
            cached_input_tokens,
            output_tokens,
            input_price_per_million,
            cached_input_price_per_million,
            output_price_per_million,
        )
    }

    /// Create a new token usage record with cache-aware input pricing.
    pub fn new_with_cache(
        model: impl Into<String>,
        input_tokens: u64,
        cached_input_tokens: u64,
        output_tokens: u64,
        input_price_per_million: f64,
        cached_input_price_per_million: f64,
        output_price_per_million: f64,
    ) -> Self {
        let model = model.into();
        let input_price_per_million = Self::sanitize_price(input_price_per_million);
        let output_price_per_million = Self::sanitize_price(output_price_per_million);
        let cached_input_price_per_million = Self::sanitize_price(cached_input_price_per_million);
        let cached_input_tokens = cached_input_tokens.min(input_tokens);
        let billable_uncached_input = input_tokens.saturating_sub(cached_input_tokens);
        let total_tokens = input_tokens.saturating_add(output_tokens);

        // Calculate cost: (tokens / 1M) * price_per_million for each band.
        // Cached subset uses its own rate when set, else falls back to the
        // standard input rate so providers without a cache-rate aren't
        // charged $0 for the cached portion.
        let cached_rate = if cached_input_price_per_million > 0.0 {
            cached_input_price_per_million
        } else {
            input_price_per_million
        };
        let input_cost = (billable_uncached_input as f64 / 1_000_000.0) * input_price_per_million;
        let cached_cost = (cached_input_tokens as f64 / 1_000_000.0) * cached_rate;
        let output_cost = (output_tokens as f64 / 1_000_000.0) * output_price_per_million;
        let cost_usd = input_cost + cached_cost + output_cost;

        Self {
            model,
            input_tokens,
            output_tokens,
            cached_input_tokens,
            total_tokens,
            cost_usd,
            pricing_available: true,
            timestamp: chrono::Utc::now(),
        }
    }

    /// Get the total cost.
    pub fn cost(&self) -> f64 {
        self.cost_usd
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum UsagePeriod {
    /// Current tracker session for this daemon lifetime.
    Session,
    /// Current UTC/local-day aggregate used by daily budget checks.
    Day,
    /// Current calendar-month aggregate used by monthly budget checks.
    Month,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostRecord {
    /// Unique identifier
    pub id: String,
    /// Token usage details
    pub usage: TokenUsage,
    /// Session identifier (for grouping)
    pub session_id: String,
    /// Alias of the agent that incurred this cost (HashMap key in
    /// `config.agents`). `None` for records persisted before per-agent
    /// attribution, or when `[cost].track_per_agent = false`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub agent_alias: Option<String>,
    #[serde(
        default,
        rename = "goal_task_id",
        alias = "task_id",
        skip_serializing_if = "Option::is_none"
    )]
    pub task_id: Option<String>,
}

impl CostRecord {
    /// Create a new cost record without agent attribution.
    pub fn new(session_id: impl Into<String>, usage: TokenUsage) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            usage,
            session_id: session_id.into(),
            agent_alias: None,
            task_id: None,
        }
    }

    /// Create a new cost record attributed to an agent.
    pub fn with_agent(
        session_id: impl Into<String>,
        agent_alias: Option<String>,
        usage: TokenUsage,
    ) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            usage,
            session_id: session_id.into(),
            agent_alias,
            task_id: None,
        }
    }

    /// Create a new cost record attributed to an agent and/or durable task.
    pub fn with_attribution(
        session_id: impl Into<String>,
        agent_alias: Option<String>,
        task_id: Option<String>,
        usage: TokenUsage,
    ) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            usage,
            session_id: session_id.into(),
            agent_alias,
            task_id,
        }
    }
}

#[derive(Debug, Clone)]
pub enum BudgetCheck {
    /// Within budget; request can proceed.
    Allowed,
    /// Warning threshold exceeded, but request can proceed.
    Warning {
        /// Current spend in the affected aggregation period before the
        /// estimated request cost is committed.
        current_usd: f64,
        /// Configured limit for the affected aggregation period.
        limit_usd: f64,
        /// Aggregation period that crossed the warning threshold.
        period: UsagePeriod,
    },
    /// Budget exceeded; request should be blocked before provider dispatch.
    Exceeded {
        /// Current spend in the affected aggregation period before the
        /// estimated request cost is committed.
        current_usd: f64,
        /// Configured limit for the affected aggregation period.
        limit_usd: f64,
        /// Aggregation period whose limit would be exceeded.
        period: UsagePeriod,
    },
}

/// Cost summary for reporting.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostSummary {
    /// Total cost for the session
    pub session_cost_usd: f64,
    /// Total cost for the day
    pub daily_cost_usd: f64,
    /// Total cost for the month
    pub monthly_cost_usd: f64,
    /// Total tokens used
    pub total_tokens: u64,
    /// Number of requests
    pub request_count: usize,
    /// Breakdown by model
    pub by_model: std::collections::HashMap<String, ModelStats>,
    /// Breakdown by agent alias. Empty when `[cost].track_per_agent =
    /// false` or when no records carry an agent_alias.
    #[serde(default)]
    pub by_agent: std::collections::HashMap<String, AgentCostStats>,
}

/// Statistics for a specific agent alias.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCostStats {
    /// Agent alias (HashMap key in `config.agents`).
    pub agent_alias: String,
    /// Total cost attributed to this agent for the period.
    pub cost_usd: f64,
    /// Total tokens attributed to this agent for the period (input + output).
    pub total_tokens: u64,
    /// Input tokens (uncached + cached). Matches each record's
    /// `input_tokens` field.
    #[serde(default)]
    pub input_tokens: u64,
    /// Output tokens.
    #[serde(default)]
    pub output_tokens: u64,
    /// Cached input tokens (subset of `input_tokens` served from the
    /// provider's prompt cache).
    #[serde(default)]
    pub cached_input_tokens: u64,
    /// Number of LLM responses attributed to this agent for the period.
    pub request_count: usize,
}

/// Statistics for a specific model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStats {
    /// Model name (upstream resource id from usage telemetry).
    pub model: String,
    /// Total cost for this model.
    pub cost_usd: f64,
    /// Total input tokens for this model
    #[serde(default)]
    pub input_tokens: u64,
    /// Total cached input tokens for this model
    #[serde(default)]
    pub cached_input_tokens: u64,
    /// Total output tokens for this model
    #[serde(default)]
    pub output_tokens: u64,
    /// Total tokens for this model
    pub total_tokens: u64,
    /// Number of LLM responses for this model.
    pub request_count: usize,
}

impl Default for CostSummary {
    fn default() -> Self {
        Self {
            session_cost_usd: 0.0,
            daily_cost_usd: 0.0,
            monthly_cost_usd: 0.0,
            total_tokens: 0,
            request_count: 0,
            by_model: std::collections::HashMap::new(),
            by_agent: std::collections::HashMap::new(),
        }
    }
}

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

    #[test]
    fn token_usage_calculation() {
        let usage = TokenUsage::new("test/model", 1000, 500, 0, 3.0, 15.0, 0.0);

        // Expected: (1000/1M)*3 + (500/1M)*15 = 0.003 + 0.0075 = 0.0105
        assert!((usage.cost_usd - 0.0105).abs() < 0.0001);
        assert_eq!(usage.input_tokens, 1000);
        assert_eq!(usage.cached_input_tokens, 0);
        assert_eq!(usage.billable_input_tokens(), 1000);
        assert_eq!(usage.output_tokens, 500);
        assert_eq!(usage.total_tokens, 1500);
    }

    #[test]
    fn token_usage_cached_input_billed_at_cached_rate() {
        // 200 cached input @ 0.5/Mtok + 800 uncached input @ 3/Mtok + 500 output @ 15/Mtok
        // = (200/1e6)*0.5 + (800/1e6)*3 + (500/1e6)*15
        // = 0.0001 + 0.0024 + 0.0075 = 0.01
        let usage = TokenUsage::new("test/model", 1000, 500, 200, 3.0, 15.0, 0.5);
        assert!((usage.cost_usd - 0.01).abs() < 1e-6, "{}", usage.cost_usd);
        assert_eq!(usage.cached_input_tokens, 200);
    }

    #[test]
    fn token_usage_zero_cached_rate_falls_back_to_input_rate() {
        // Cached rate 0 means "no discount" — bill cached subset at the
        // standard input rate so providers without a published cache rate
        // still produce a sane total.
        let with_cache = TokenUsage::new("test/model", 1000, 500, 200, 3.0, 15.0, 0.0);
        let without_cache = TokenUsage::new("test/model", 1000, 500, 0, 3.0, 15.0, 0.0);
        assert!((with_cache.cost_usd - without_cache.cost_usd).abs() < 1e-9);
    }

    #[test]
    fn token_usage_cache_aware_calculation() {
        let usage = TokenUsage::new_with_cache("test/model", 1_000, 800, 500, 3.0, 0.3, 15.0);

        // Expected: uncached=(200/1M)*3 + cached=(800/1M)*0.3 + output=(500/1M)*15
        let expected = 0.0006 + 0.00024 + 0.0075;
        assert!((usage.cost_usd - expected).abs() < 0.0001);
        assert_eq!(usage.cached_input_tokens, 800);
        assert_eq!(usage.billable_input_tokens(), 200);
        assert_eq!(usage.total_tokens, 1500);
    }

    #[test]
    fn token_usage_zero_tokens() {
        let usage = TokenUsage::new("test/model", 0, 0, 0, 3.0, 15.0, 0.0);
        assert!(usage.cost_usd.abs() < f64::EPSILON);
        assert_eq!(usage.total_tokens, 0);
    }

    #[test]
    fn token_usage_negative_or_non_finite_prices_are_clamped() {
        let usage = TokenUsage::new("test/model", 1000, 1000, 0, -3.0, f64::NAN, f64::INFINITY);
        assert!(usage.cost_usd.abs() < f64::EPSILON);
        assert_eq!(usage.total_tokens, 2000);
    }

    #[test]
    fn cost_record_creation() {
        let usage = TokenUsage::new("test/model", 100, 50, 0, 1.0, 2.0, 0.0);
        let record = CostRecord::new("session-123", usage);

        assert_eq!(record.session_id, "session-123");
        assert!(record.task_id.is_none());
        assert!(!record.id.is_empty());
        assert_eq!(record.usage.model, "test/model");
    }

    #[test]
    fn cost_record_task_attribution_roundtrips_existing_field() {
        let usage = TokenUsage::new("test/model", 100, 50, 0, 1.0, 2.0, 0.0);
        let record = CostRecord::with_attribution(
            "session-123",
            Some("agent-a".into()),
            Some("goal-123".into()),
            usage,
        );
        let json = serde_json::to_string(&record).unwrap();
        assert!(
            json.contains("goal_task_id"),
            "serialized key stays compatible with existing JSONL"
        );
        let parsed: CostRecord = serde_json::from_str(&json).unwrap();

        assert_eq!(parsed.agent_alias.as_deref(), Some("agent-a"));
        assert_eq!(parsed.task_id.as_deref(), Some("goal-123"));
    }

    #[test]
    fn cost_record_task_attribution_accepts_generic_alias() {
        let json = r#"{
            "id": "rec-1",
            "usage": {
                "model": "test/model",
                "input_tokens": 10,
                "output_tokens": 5,
                "total_tokens": 15,
                "cost_usd": 0.001,
                "timestamp": "2026-06-30T00:00:00Z"
            },
            "session_id": "session-123",
            "agent_alias": "agent-a",
            "task_id": "task-123"
        }"#;
        let parsed: CostRecord = serde_json::from_str(json).unwrap();

        assert_eq!(parsed.task_id.as_deref(), Some("task-123"));
        assert!(parsed.usage.pricing_available);
    }
}
