cortex-rs-stats 0.2.0

use anyhow::Result;
use chrono::{Duration, Utc};
use std::collections::HashMap;

use crate::{ModelUsage, StatsSource, UsageSummary};

pub async fn fetch_anthropic(admin_key: &str, days: u32) -> Result<UsageSummary> {
    let client = reqwest::Client::new();

    let end = Utc::now();
    let start = end - Duration::days(days as i64);

    let url = "https://api.anthropic.com/v1/organizations/usage_report/messages";
    let response = client
        .get(url)
        .header("x-api-key", admin_key)
        .header("anthropic-version", "2023-06-01")
        .query(&[
            ("starting_at", start.format("%Y-%m-%d").to_string()),
            ("ending_at", end.format("%Y-%m-%d").to_string()),
            ("bucket_width", "1d".to_string()),
        ])
        .send()
        .await?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("Anthropic API returned {}: {}", status, body);
    }

    let data: serde_json::Value = response.json().await?;
    let by_model = parse_anthropic_response(&data);
    let total = by_model.iter().map(|m| m.cost_usd).sum::<f64>();

    Ok(UsageSummary {
        days,
        by_model,
        total_cost_usd: total,
        routing_savings_usd: 0.0,
        source: StatsSource::ProviderApi,
    })
}

pub async fn fetch_openai(admin_key: &str, days: u32) -> Result<UsageSummary> {
    let client = reqwest::Client::new();
    let end = Utc::now();
    let start = end - Duration::days(days as i64);

    let url = "https://api.openai.com/v1/organization/usage/completions";
    let response = client
        .get(url)
        .bearer_auth(admin_key)
        .query(&[
            ("start_time", start.timestamp().to_string()),
            ("end_time", end.timestamp().to_string()),
            ("bucket_width", "1d".to_string()),
            ("group_by", "model".to_string()),
        ])
        .send()
        .await?;

    if !response.status().is_success() {
        let status = response.status();
        let body = response.text().await.unwrap_or_default();
        anyhow::bail!("OpenAI API returned {}: {}", status, body);
    }

    let data: serde_json::Value = response.json().await?;
    let by_model = parse_openai_response(&data);
    let total = by_model.iter().map(|m| m.cost_usd).sum::<f64>();

    Ok(UsageSummary {
        days,
        by_model,
        total_cost_usd: total,
        routing_savings_usd: 0.0,
        source: StatsSource::ProviderApi,
    })
}

fn parse_openai_response(data: &serde_json::Value) -> Vec<ModelUsage> {
    // OpenAI usage API returns { data: [{ start_time, end_time, results: [{ model, input_tokens, output_tokens, ... }] }] }
    // Cost requires a second call to /organization/costs; for v0.2 we estimate from token counts.
    let mut by_model: HashMap<String, (u64, u64)> = HashMap::new();
    if let Some(buckets) = data.get("data").and_then(|d| d.as_array()) {
        for bucket in buckets {
            if let Some(results) = bucket.get("results").and_then(|r| r.as_array()) {
                for entry in results {
                    let model = entry
                        .get("model")
                        .and_then(|m| m.as_str())
                        .unwrap_or("openai-unknown")
                        .to_string();
                    let input = entry
                        .get("input_tokens")
                        .and_then(|t| t.as_u64())
                        .unwrap_or(0);
                    let output = entry
                        .get("output_tokens")
                        .and_then(|t| t.as_u64())
                        .unwrap_or(0);
                    let e = by_model.entry(model).or_default();
                    e.0 += input;
                    e.1 += output;
                }
            }
        }
    }

    // OpenAI API doesn't return cost directly here; pricing is applied by the caller's config.
    // Return entries with cost_usd=0; the caller's pricing table merges them.
    by_model
        .into_iter()
        .map(|(model, (input, output))| ModelUsage {
            model,
            input_tokens: input,
            output_tokens: output,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            cost_usd: 0.0,
            pct: 0.0,
        })
        .collect()
}

fn parse_anthropic_response(data: &serde_json::Value) -> Vec<ModelUsage> {
    let mut by_model: HashMap<String, (u64, u64, f64)> = HashMap::new();

    if let Some(buckets) = data.get("buckets").and_then(|b| b.as_array()) {
        for bucket in buckets {
            if let Some(models) = bucket.get("models").and_then(|m| m.as_array()) {
                for entry in models {
                    let model = entry
                        .get("model")
                        .and_then(|m| m.as_str())
                        .unwrap_or("unknown")
                        .to_string();
                    let input = entry
                        .get("input_tokens")
                        .and_then(|t| t.as_u64())
                        .unwrap_or(0);
                    let output = entry
                        .get("output_tokens")
                        .and_then(|t| t.as_u64())
                        .unwrap_or(0);
                    // cost is returned in USD cents as decimal string
                    let cost_cents: f64 = entry
                        .get("cost")
                        .and_then(|c| c.as_str())
                        .and_then(|s| s.parse().ok())
                        .unwrap_or(0.0);

                    let e = by_model.entry(model).or_default();
                    e.0 += input;
                    e.1 += output;
                    e.2 += cost_cents / 100.0; // cents → USD
                }
            }
        }
    }

    let total_cost: f64 = by_model.values().map(|(_, _, c)| c).sum();

    let mut usage: Vec<ModelUsage> = by_model
        .into_iter()
        .map(|(model, (input, output, cost))| ModelUsage {
            model,
            input_tokens: input,
            output_tokens: output,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            cost_usd: cost,
            pct: if total_cost > 0.0 { cost / total_cost * 100.0 } else { 0.0 },
        })
        .collect();

    usage.sort_by(|a, b| b.cost_usd.partial_cmp(&a.cost_usd).unwrap_or(std::cmp::Ordering::Equal));
    usage
}