//! Budgets: monthly limits tracked from recorded costs, plus the real balance
//! for providers that expose one (OpenRouter, fal.ai), and price discovery.

use crate::app::App;
use crate::config::{ProviderConfig, expand_env};
use anyhow::{Context, Result, bail};
use chrono::{Datelike, TimeZone};
use serde_json::{Value, json};
use std::time::Duration;

/// Start of the current month in local time (ms).
pub fn month_start() -> i64 {
    let now = chrono::Local::now();
    chrono::Local
        .with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
        .single()
        .map(|d| d.timestamp_millis())
        .unwrap_or(0)
}

fn is_openrouter(p: &ProviderConfig) -> bool {
    p.base_url.contains("openrouter.ai")
}

/// Month-to-date cost from Anthropic's cost report (organisation-wide, USD).
async fn anthropic_cost(p: &ProviderConfig) -> Result<Value> {
    let admin = expand_env(&p.admin_key);
    let start = chrono::DateTime::from_timestamp_millis(month_start()).unwrap_or_default();
    let end = chrono::Utc::now() + chrono::Duration::days(1);
    let base = if p.base_url.is_empty() { "https://api.anthropic.com/v1".to_string() } else { expand_env(&p.base_url).trim_end_matches('/').to_string() };
    let mut total_cents = 0.0;
    let mut page: Option<String> = None;
    for _ in 0..10 {
        let mut url = reqwest::Url::parse(&format!("{base}/organizations/cost_report"))?;
        url.query_pairs_mut()
            .append_pair("starting_at", &start.format("%Y-%m-%dT00:00:00Z").to_string())
            .append_pair("ending_at", &end.format("%Y-%m-%dT00:00:00Z").to_string())
            .append_pair("limit", "31");
        if let Some(pg) = &page {
            url.query_pairs_mut().append_pair("page", pg);
        }
        let resp = crate::provider::http_client()
            .get(url)
            .header("x-api-key", &admin)
            .header("anthropic-version", "2026-10-21")
            .timeout(Duration::from_secs(20))
            .send()
            .await?;
        if !resp.status().is_success() {
            return Err(crate::provider::http_error(resp).await);
        }
        let v: Value = resp.json().await?;
        for bucket in v["data"].as_array().cloned().unwrap_or_default() {
            for r in bucket["results"].as_array().cloned().unwrap_or_default() {
                // Amounts are decimal strings in cents.
                total_cents += r["amount"].as_str().and_then(|a| a.parse::<f64>().ok()).or_else(|| r["amount"].as_f64()).unwrap_or(0.0);
            }
        }
        if v["has_more"].as_bool() != Some(true) {
            break;
        }
        page = v["next_page"].as_str().map(String::from);
    }
    Ok(json!({"currency": "USD", "spent_month": total_cents / 100.0, "scope": "organization"}))
}

/// Balance or cost as reported by the provider, if it has an API for it.
async fn remote_balance(p: &ProviderConfig) -> Option<Result<Value>> {
    if p.kind == crate::config::ProviderKind::Anthropic && !expand_env(&p.admin_key).is_empty() {
        return Some(anthropic_cost(p).await);
    }
    if !is_openrouter(p) {
        return None;
    }
    Some(
        async {
            let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
            let resp = crate::provider::http_client()
                .get(format!("{base}/key"))
                .bearer_auth(expand_env(&p.api_key))
                .timeout(Duration::from_secs(15))
                .send()
                .await?
                .error_for_status()?;
            let v: Value = resp.json().await?;
            let d = &v["data"];
            Ok(json!({
                "currency": "USD",
                "limit": d["limit"],
                "used": d["usage"],
                "remaining": d.get("limit_remaining").cloned().unwrap_or_else(|| match (d["limit"].as_f64(), d["usage"].as_f64()) {
                    (Some(l), Some(u)) => json!(l - u),
                    _ => Value::Null,
                }),
            }))
        }
        .await,
    )
}

async fn fal_balance(admin_key: &str) -> Result<Value> {
    let resp = crate::provider::http_client()
        .get("https://api.fal.ai/v1/account/billing?expand=credits")
        .header("Authorization", format!("Key {admin_key}"))
        .timeout(Duration::from_secs(15))
        .send()
        .await?;
    if !resp.status().is_success() {
        return Err(crate::provider::http_error(resp).await);
    }
    let v: Value = resp.json().await?;
    Ok(json!({
        "currency": v["credits"]["currency"],
        "remaining": v["credits"]["current_balance"],
    }))
}

fn entry(name: &str, budget: Option<f64>, spent: f64, warn: f64) -> Value {
    let share = budget.filter(|b| *b > 0.0).map(|b| spent / b);
    json!({
        "provider": name,
        "budget": budget,
        "spent": spent,
        "remaining": budget.map(|b| b - spent),
        "share": share,
        "warning": share.is_some_and(|s| s >= warn),
        "exceeded": share.is_some_and(|s| s >= 1.0),
    })
}

/// Budget overview for the current month.
/// Subscription windows of opencode Go, as share of the plan's monthly limit: rolling
/// 5 hours (20 %) and 7 days (50 %), and the calendar month in UTC (100 %). Only usage
/// through CodeCraft is counted.
pub fn plan_usage(app: &App, provider: &str, limit: f64) -> Result<Value> {
    let now = chrono::Utc::now();
    let month = chrono::Utc.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0).single().map(|d| d.timestamp_millis()).unwrap_or(0);
    let ms = now.timestamp_millis();
    let window = |name: &str, from: i64, share: f64| -> Result<Value> {
        let used = app.store.allowance_since(provider, from)?;
        let cap = limit * share;
        Ok(json!({"window": name, "used": used, "limit": cap, "ratio": if cap > 0.0 { used / cap } else { 0.0 }}))
    };
    Ok(json!({
        "limit": limit,
        "windows": [
            window("5h", ms - 5 * 3_600_000, 0.2)?,
            window("week", ms - 7 * 86_400_000, 0.5)?,
            window("month", month, 1.0)?,
        ],
    }))
}

pub async fn overview(app: &App) -> Result<Value> {
    let cfg = app.config();
    let spent = app.store.cost_by_provider(month_start())?;
    let warn = cfg.server.budget_warning;
    let mut rows = Vec::new();
    let mut tasks = Vec::new();
    for (name, p) in &cfg.providers {
        if !p.enabled {
            continue;
        }
        let mut row = entry(name, p.monthly_budget, spent.get(name).copied().unwrap_or(0.0), warn);
        if let Some(limit) = p.plan_limit.filter(|l| *l > 0.0) {
            row["plan"] = plan_usage(app, name, limit)?;
        }
        rows.push(row);
        let p = p.clone();
        tasks.push(tokio::spawn(async move { remote_balance(&p).await }));
    }
    for (row, task) in rows.iter_mut().zip(tasks) {
        if let Ok(Some(res)) = task.await {
            row["remote"] = match res {
                Ok(v) => v,
                Err(e) => json!({"error": format!("{e:#}")}),
            };
        }
    }
    if cfg.images.enabled {
        let mut row = entry("fal", cfg.images.monthly_budget, spent.get("fal").copied().unwrap_or(0.0), warn);
        let admin = expand_env(&cfg.images.admin_key);
        if !admin.is_empty() {
            row["remote"] = match fal_balance(&admin).await {
                Ok(v) => v,
                Err(e) => json!({"error": format!("{e:#}")}),
            };
        }
        row["images"] = json!(true);
        rows.push(row);
    }
    let total: f64 = spent.values().sum();
    let global = (cfg.server.monthly_budget > 0.0).then_some(cfg.server.monthly_budget);
    Ok(json!({
        "month_start": month_start(),
        "currency": cfg.server.currency,
        "total": entry("total", global, total, warn),
        "providers": rows,
    }))
}

/// Per-token prices from the provider's model list, converted to per 1M tokens.
fn is_opencode(p: &ProviderConfig) -> bool {
    p.base_url.contains("opencode.ai")
}

async fn get_json(url: &str) -> Result<Value> {
    Ok(crate::provider::http_client().get(url).timeout(Duration::from_secs(30)).send().await?.error_for_status()?.json().await?)
}

/// One opencode condition object → (days, hours, from, until); None if it uses
/// something we do not understand.
type When = (Vec<String>, Vec<(String, String)>, Option<String>, Option<String>);

fn opencode_when(w: &Value) -> Option<When> {
    let obj = w.as_object()?;
    let (mut days, mut hours, mut from, mut until) = (Vec::new(), Vec::new(), None, None);
    // A pair of utc_hours bounds, given as {"all": [{"utc_hours": {"gte"}}, {"utc_hours": {"lt"}}]}
    // or directly as {"utc_hours": {"gte", "lt"}}.
    let range = |v: &Value| -> Option<(String, String)> {
        if let Some(h) = v.get("utc_hours") {
            return Some((h["gte"].as_str()?.to_string(), h["lt"].as_str()?.to_string()));
        }
        let all = v.get("all")?.as_array()?;
        let get = |k: &str| all.iter().find_map(|x| x["utc_hours"][k].as_str()).map(String::from);
        Some((get("gte")?, get("lt")?))
    };
    for (k, v) in obj {
        match k.as_str() {
            "utc_days" => days = v["in"].as_array()?.iter().map(|d| d.as_str().map(String::from)).collect::<Option<_>>()?,
            "any" => hours = v.as_array()?.iter().map(range).collect::<Option<_>>()?,
            "all" | "utc_hours" => hours = vec![range(w)?],
            "date" => {
                from = v["gte"].as_str().map(String::from);
                until = v["lt"].as_str().map(String::from);
                if from.is_none() && until.is_none() {
                    return None;
                }
            }
            _ => return None,
        }
    }
    Some((days, hours, from, until))
}

/// Model row from opencode's model list (prices per 1M tokens, tiers, time factors).
fn opencode_model(m: &Value) -> Value {
    let cost = &m["cost"]["text"];
    let (input, output) = (&cost["input"], &cost["output"]);
    // Context tiers: input and output overrides share their thresholds.
    let mut tiers: std::collections::BTreeMap<u64, serde_json::Map<String, Value>> = Default::default();
    for (side, field) in [(input, "input_price"), (output, "output_price")] {
        for o in side["overrides"].as_array().into_iter().flatten() {
            let Some(gt) = o["when"]["context"]["gt"].as_u64() else { continue };
            let t = tiers.entry(gt).or_default();
            t.insert(field.into(), o["price"].clone());
            if field == "input_price" {
                if o["cache"]["read"].is_number() {
                    t.insert("cache_read_price".into(), o["cache"]["read"].clone());
                }
                if o["cache"]["write"].is_number() {
                    t.insert("cache_write_price".into(), o["cache"]["write"].clone());
                }
            }
        }
    }
    let price_tiers: Vec<Value> = tiers.into_iter().map(|(gt, mut t)| {
        t.insert("above_context".into(), json!(gt));
        Value::Object(t)
    }).collect();
    let mut notes = Vec::new();
    let price_multipliers: Vec<Value> = m["cost"]["multipliers"]
        .as_array()
        .into_iter()
        .flatten()
        .filter_map(|x| match opencode_when(&x["when"]) {
            Some((days, hours, from, until)) => Some(json!({"factor": x["factor"], "days": days, "hours_utc": hours, "from": from, "until": until})),
            None => {
                notes.push(format!("unsupported price condition: {}", x["when"]));
                None
            }
        })
        .map(|mut v| {
            // Drop empty conditions so the config stays short.
            v.as_object_mut().unwrap().retain(|_, x| !(x.is_null() || x.as_array().is_some_and(|a| a.is_empty())));
            v
        })
        .collect();
    let images = m["modalities"]["input"].as_array().is_some_and(|a| a.iter().any(|x| x.as_str().is_some_and(|s| s.starts_with("image"))));
    json!({
        "id": m["id"],
        "context": m["limit"]["context"],
        "max_output": m["limit"]["output"],
        "vision": images,
        "reasoning": m["capabilities"]["reasoning"]["supported"].as_bool().unwrap_or(false),
        "input_price": input["price"],
        "output_price": output["price"],
        "cache_read_price": input["cache"]["read"],
        "cache_write_price": input["cache"]["write"],
        "price_tiers": price_tiers,
        "price_multipliers": price_multipliers,
        "notes": notes,
    })
}

/// opencode (Zen and Go) publishes prices in its model list. The Anthropic-protocol
/// endpoint lists models without them, so fall back to the gateway's own /v1/models.
async fn opencode_prices(base: &str) -> Result<Value> {
    let mut urls = vec![format!("{base}/models")];
    for suffix in ["/openai/v1", "/anthropic/v1"] {
        if let Some(root) = base.strip_suffix(suffix) {
            urls.push(format!("{root}/v1/models"));
        }
    }
    for url in &urls {
        let Ok(v) = get_json(url).await else { continue };
        let models: Vec<Value> = v["data"].as_array().into_iter().flatten().filter(|m| m.get("cost").is_some()).map(opencode_model).collect();
        if !models.is_empty() {
            return Ok(json!({"currency": "USD", "models": models}));
        }
    }
    bail!("opencode did not return prices for this endpoint")
}

pub async fn discover_prices(app: &App, provider: &str) -> Result<Value> {
    let cfg = app.config();
    let p = cfg.providers.get(provider).context("unknown provider")?;
    let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
    if is_opencode(p) {
        return opencode_prices(&base).await;
    }
    if !is_openrouter(p) {
        bail!("this provider does not publish prices via its API – enter them manually");
    }
    let v: Value = crate::provider::http_client()
        .get(format!("{base}/models"))
        .timeout(Duration::from_secs(30))
        .send()
        .await?
        .error_for_status()?
        .json()
        .await?;
    let per_m = |x: &Value| x.as_str().and_then(|s| s.parse::<f64>().ok()).or_else(|| x.as_f64()).map(|f| (f * 1_000_000.0 * 10_000.0).round() / 10_000.0);
    let list: Vec<Value> = v["data"]
        .as_array()
        .cloned()
        .unwrap_or_default()
        .iter()
        .map(|m| {
            json!({
                "id": m["id"],
                "context": m["context_length"],
                "input_price": per_m(&m["pricing"]["prompt"]),
                "output_price": per_m(&m["pricing"]["completion"]),
                "cache_read_price": per_m(&m["pricing"]["input_cache_read"]),
                "cache_write_price": per_m(&m["pricing"]["input_cache_write"]),
            })
        })
        .collect();
    Ok(json!({"currency": "USD", "models": list}))
}

#[cfg(test)]
mod tests {
    #[test]
    fn opencode_cost_conversion() {
        let m: serde_json::Value = serde_json::from_str(r#"{"id":"grok","limit":{"context":500000,"output":8000},
            "modalities":{"input":["text","image/png"]},"capabilities":{"reasoning":{"supported":true}},
            "cost":{"text":{"input":{"price":2,"cache":{"read":0.5},"overrides":[{"when":{"context":{"gt":200000}},"price":4,"cache":{"read":1}}]},
                            "output":{"price":6,"overrides":[{"when":{"context":{"gt":200000}},"price":12}]}},
                    "multipliers":[{"when":{"utc_days":{"in":["mon","tue"]},"any":[{"all":[{"utc_hours":{"gte":"01:00"}},{"utc_hours":{"lt":"04:00"}}]}]},"factor":2},
                                   {"when":{"date":{"gte":"2026-12-01"}},"factor":2},
                                   {"when":{"region":"eu"},"factor":3}]}}"#).unwrap();
        let r = super::opencode_model(&m);
        assert_eq!(r["input_price"], 2);
        assert_eq!(r["cache_read_price"], 0.5);
        assert_eq!(r["vision"], true);
        assert_eq!(r["price_tiers"][0], serde_json::json!({"above_context": 200000, "input_price": 4, "cache_read_price": 1, "output_price": 12}));
        assert_eq!(r["price_multipliers"][0], serde_json::json!({"factor": 2, "days": ["mon", "tue"], "hours_utc": [["01:00", "04:00"]]}));
        assert_eq!(r["price_multipliers"][1], serde_json::json!({"factor": 2, "from": "2026-12-01"}));
        assert_eq!(r["notes"].as_array().unwrap().len(), 1, "unknown condition is reported, not guessed");
        // The result must deserialize into the model config.
        let mc: crate::config::ModelConfig = serde_json::from_value(serde_json::json!({"id": "grok", "price_tiers": r["price_tiers"], "price_multipliers": r["price_multipliers"]})).unwrap();
        assert_eq!(mc.price_tiers[0].above_context, 200000);
    }

    #[test]
    fn budget_entry_flags() {
        let e = super::entry("x", Some(10.0), 8.5, 0.8);
        assert_eq!(e["warning"], true);
        assert_eq!(e["exceeded"], false);
        assert_eq!(e["remaining"], 1.5);
        let none = super::entry("y", None, 3.0, 0.8);
        assert_eq!(none["warning"], false);
    }
}

/// Refuse a request when a budget or the owner's quota is used up.
/// `provider` is the provider name ("fal" for images).
pub fn check(app: &App, provider: &str, owner_id: Option<&str>) -> Result<()> {
    let cfg = app.config();
    let from = month_start();
    if let Some(user) = owner_id.and_then(|id| app.store.user(id).ok().flatten()) {
        if let Some(q) = user.monthly_quota.filter(|q| *q > 0.0) {
            let spent = app.store.user_cost_since(&user.id, from)?;
            if spent >= q {
                bail!(
                    "monthly quota of {q:.2} {cur} for user {} is used up ({spent:.2} {cur} spent) – ask an admin to raise it",
                    user.username,
                    cur = cfg.server.currency
                );
            }
        }
    }
    if cfg.server.budget_action != "block" {
        return Ok(());
    }
    let spent = app.store.cost_by_provider(from)?;
    let total: f64 = spent.values().sum();
    if cfg.server.monthly_budget > 0.0 && total >= cfg.server.monthly_budget {
        bail!("the monthly budget of {:.2} {} is used up", cfg.server.monthly_budget, cfg.server.currency);
    }
    let limit = if provider == "fal" { cfg.images.monthly_budget } else { cfg.providers.get(provider).and_then(|p| p.monthly_budget) };
    if let Some(b) = limit.filter(|b| *b > 0.0) {
        let s = spent.get(provider).copied().unwrap_or(0.0);
        if s >= b {
            bail!("the monthly budget of {b:.2} {} for {provider} is used up", cfg.server.currency);
        }
    }
    Ok(())
}