quotch 0.5.1

use std::time::Duration;

use chrono::{DateTime, TimeZone, Utc};
use serde_json::Value;

use crate::model::{Account, CredentialSource, Snapshot, Status, Unit, Window, WindowKind};
use crate::providers::{FetchError, Provider};

// ponytail: Antigravity's PUBLIC embedded installed-app OAuth client, baked into
// the IDE binary — NOT a real secret. These mirror Antigravity's private
// internals and may drift if Google rotates the shipped client. Refreshing an
// installed-app grant requires the client_secret alongside the id.
const CLIENT_ID: &str = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com";
const CLIENT_SECRET: &str = "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";
const TOKEN_URL: &str = "https://oauth2.googleapis.com/token";

// ponytail: Google's Cloud Code backend routes account TIER and quota by the
// caller's client identity — a "quotch/x" UA resolves to a different tier whose
// retrieveUserQuotaSummary 403s (verified live: Antigravity UA → 200 with quota
// buckets; quotch UA → 403). We read our own account's quota through the client
// Google expects, so we must present Antigravity's UA on every call. Bump the
// version if Antigravity's shipped UA drifts and the endpoint starts 403ing.
const ANTIGRAVITY_UA: &str = "vscode/1.100.0 (Antigravity/1.107.0)";

// Cloud Code backends, tried in order. Advance to the next only on HTTP 429 or
// 5xx (and transport failures); any other completed response is used as-is.
const HOSTS: [&str; 3] = [
    "https://daily-cloudcode-pa.sandbox.googleapis.com",
    "https://daily-cloudcode-pa.googleapis.com",
    "https://cloudcode-pa.googleapis.com",
];

pub struct Antigravity;

impl Provider for Antigravity {
    fn id(&self) -> &'static str {
        "antigravity"
    }

    fn discover(&self) -> Vec<Account> {
        // Emit the default account only when the agy CLI's OAuth token file
        // exists. A wholly absent file means the user doesn't use Antigravity at
        // all, so we stay silent rather than nag (mirror copilot). A
        // present-but-broken file still surfaces as auth_missing via fetch.
        // Offline only: no network, no token read here.
        let path = std::env::home_dir()
            .unwrap_or_default()
            .join(".gemini")
            .join("antigravity-cli")
            .join("antigravity-oauth-token");
        if !path.exists() {
            return vec![];
        }
        vec![Account {
            provider: "antigravity",
            id: "default".into(),
            source: CredentialSource::FilePath(path),
            label: None,
            display: None,
        }]
    }

    fn fetch(&self, acct: &Account) -> Result<Snapshot, FetchError> {
        let path = match &acct.source {
            CredentialSource::FilePath(p) => p,
            _ => return Err(FetchError::AuthMissing),
        };

        // Missing file / unreadable / bad JSON all collapse to AuthMissing.
        let creds: Value = std::fs::read_to_string(path)
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok())
            .ok_or(FetchError::AuthMissing)?;

        let refresh = creds["token"]["refresh_token"]
            .as_str()
            .ok_or(FetchError::AuthMissing)?;

        // Mint a fresh access token; never persisted, used only in headers below.
        let access = refresh_token(refresh)?;

        let agent = build_agent();

        // Plan lookup is non-fatal: on any failure both come back None and we
        // still report quota. `host` is the backend that answered so the quota
        // call can prefer it.
        let (project_id, plan, host) = load_code_assist(&agent, &access);

        let summary = retrieve_quota_summary(&agent, &access, project_id.as_deref(), host)?;

        Ok(Snapshot {
            provider: "antigravity".into(),
            account: acct.id.clone(),
            label: acct.label.clone(),
            plan,
            windows: parse_quota_summary(&summary),
            fetched_at: Utc::now(),
            status: Status::Ok,
            error: None,
            raw: Some(summary),
        })
    }
}

// ponytail: same 3s connect / 8s overall split as copilot.rs / claude.rs — a
// cold TLS handshake on a slow link can exceed a tighter overall timeout and pin
// the cache stale. Warm runs never pay this; speed lives in the cache path.
fn build_agent() -> ureq::Agent {
    ureq::AgentBuilder::new()
        .timeout_connect(Duration::from_secs(3))
        .timeout(Duration::from_secs(8))
        .build()
}

// Exchange the stored refresh token for a short-lived access token. Non-2xx or
// an `invalid_grant` body means the grant is dead → AuthMissing; a transport
// failure is Network. The token is never logged or stored.
fn refresh_token(refresh: &str) -> Result<String, FetchError> {
    let agent = build_agent();
    let (ok, body) = match agent
        .post(TOKEN_URL)
        .set("User-Agent", ANTIGRAVITY_UA)
        .send_form(&[
            ("client_id", CLIENT_ID),
            ("client_secret", CLIENT_SECRET),
            ("refresh_token", refresh),
            ("grant_type", "refresh_token"),
        ]) {
        Ok(r) => (
            true,
            r.into_string()
                .map_err(|e| FetchError::Network(e.to_string()))?,
        ),
        Err(ureq::Error::Status(_, r)) => (false, r.into_string().unwrap_or_default()),
        Err(ureq::Error::Transport(t)) => return Err(FetchError::Network(t.to_string())),
    };

    if !ok || body.contains("invalid_grant") {
        return Err(FetchError::AuthMissing);
    }
    serde_json::from_str::<Value>(&body)
        .ok()
        .and_then(|j| j["access_token"].as_str().map(str::to_string))
        .ok_or(FetchError::AuthMissing)
}

// One bearer-authed JSON POST. Returns the completed HTTP status and body for any
// status (ureq surfaces non-2xx as Err; both arms normalize to the same shape).
// None marks a transport failure so callers advance to the next host.
fn post_json(agent: &ureq::Agent, url: &str, token: &str, body: &Value) -> Option<(u16, String)> {
    let req = agent
        .post(url)
        .set("Authorization", &format!("Bearer {token}"))
        .set("Content-Type", "application/json")
        .set("User-Agent", ANTIGRAVITY_UA);
    match req.send_string(&body.to_string()) {
        Ok(r) => {
            let code = r.status();
            let s = r.into_string().ok()?;
            Some((code, s))
        }
        Err(ureq::Error::Status(code, r)) => Some((code, r.into_string().unwrap_or_default())),
        Err(ureq::Error::Transport(_)) => None,
    }
}

// POST `body` across `hosts` in order, advancing past transport failures, 429 and
// 5xx. Returns the first host to yield any other status, with its code and body;
// None if every host was exhausted.
fn sweep<'a>(
    agent: &ureq::Agent,
    hosts: &[&'a str],
    path: &str,
    token: &str,
    body: &Value,
) -> Option<(&'a str, u16, String)> {
    for &host in hosts {
        match post_json(agent, &format!("{host}{path}"), token, body) {
            Some((code, _)) if code == 429 || (500..600).contains(&code) => continue,
            Some((code, s)) => return Some((host, code, s)),
            None => continue,
        }
    }
    None
}

// Step 4: plan + project id. Non-fatal — any failure yields all-None so quota
// still reports. On a 200 it also returns the answering host, so the quota call
// can prefer it.
fn load_code_assist(
    agent: &ureq::Agent,
    token: &str,
) -> (Option<String>, Option<String>, Option<&'static str>) {
    let body = serde_json::json!({ "metadata": { "ideType": "ANTIGRAVITY" } });
    match sweep(agent, &HOSTS, "/v1internal:loadCodeAssist", token, &body) {
        Some((host, 200, s)) => {
            let json: Value = serde_json::from_str(&s).unwrap_or(Value::Null);
            (
                project_id_from_load(&json),
                plan_from_load(&json),
                Some(host),
            )
        }
        _ => (None, None, None),
    }
}

// cloudaicompanionProject may be a bare id string or an object carrying `.id`.
fn project_id_from_load(v: &Value) -> Option<String> {
    let p = &v["cloudaicompanionProject"];
    if let Some(s) = p.as_str() {
        return Some(s.to_string());
    }
    p["id"].as_str().map(str::to_string)
}

// Deliberately currentTier, NOT paidTier: paidTier is the *eligible* upsell tier
// and would mislabel a free account as paid.
fn plan_from_load(v: &Value) -> Option<String> {
    let tier = &v["currentTier"];
    tier["name"]
        .as_str()
        .or_else(|| tier["id"].as_str())
        .map(str::to_string)
}

// Step 5: primary quota. Prefer the host that answered step 4, else sweep the
// full list. A 403 while scoped to a project gets one blind retry with `{}`.
fn retrieve_quota_summary(
    agent: &ureq::Agent,
    token: &str,
    project_id: Option<&str>,
    preferred: Option<&'static str>,
) -> Result<Value, FetchError> {
    let hosts = ordered_hosts(preferred);
    let path = "/v1internal:retrieveUserQuotaSummary";
    let body = match project_id {
        Some(p) => serde_json::json!({ "project": p }),
        None => serde_json::json!({}),
    };

    let (host, code, body_str) = sweep(agent, &hosts, path, token, &body)
        .ok_or_else(|| FetchError::Network("all quota hosts exhausted".into()))?;

    match code {
        401 => Err(FetchError::AuthMissing),
        403 if project_id.is_some() => {
            // The project scoping may itself be the problem — retry once, unscoped.
            match post_json(
                agent,
                &format!("{host}{path}"),
                token,
                &serde_json::json!({}),
            ) {
                Some((200..=299, s)) => parse_summary(&s),
                _ => Err(FetchError::AuthMissing),
            }
        }
        403 => Err(FetchError::AuthMissing),
        200..=299 => parse_summary(&body_str),
        other => Err(FetchError::Parse(format!("http {other}"))),
    }
}

// Host order for step 5: the step-4 winner first (deduped), then the defaults.
fn ordered_hosts(preferred: Option<&'static str>) -> Vec<&'static str> {
    match preferred {
        Some(h) => {
            let mut v = vec![h];
            v.extend(HOSTS.iter().copied().filter(|&x| x != h));
            v
        }
        None => HOSTS.to_vec(),
    }
}

fn parse_summary(body: &str) -> Result<Value, FetchError> {
    serde_json::from_str(body).map_err(|e| FetchError::Parse(e.to_string()))
}

// PURE: turn the quota summary into windows. Iterates groups[].buckets[]; a
// bucket without remainingFraction has no usage signal and is skipped rather
// than fabricated. Deterministic: sorted by key.
fn parse_quota_summary(raw: &Value) -> Vec<Window> {
    let Some(groups) = raw.get("groups").and_then(Value::as_array) else {
        return vec![];
    };

    let mut windows: Vec<Window> = Vec::new();
    for group in groups {
        let Some(buckets) = group.get("buckets").and_then(Value::as_array) else {
            continue;
        };
        for bucket in buckets {
            // No remainingFraction → no usage signal; skip (mirror copilot).
            let Some(frac) = bucket["remainingFraction"].as_f64() else {
                continue;
            };
            let used_pct = (1.0 - frac) * 100.0;

            let bucket_id = bucket["bucketId"].as_str().unwrap_or_default();
            // Drop the trailing window suffix: "gemini-5h" → "gemini",
            // "3p-weekly" → "3p"; no '-' → whole id.
            let slug = bucket_id
                .rsplit_once('-')
                .map_or(bucket_id, |(head, _)| head);

            let key = match bucket["window"].as_str().unwrap_or_default() {
                "5h" => format!("5h:{slug}"),
                "weekly" => format!("7d:{slug}"),
                // Never drop a window on an unrecognized cadence.
                other => format!("{other}:{slug}"),
            };

            windows.push(Window {
                key,
                // Verified: both cadences refresh rolling-from-use.
                kind: WindowKind::Rolling,
                unit: Unit::Percent,
                used_pct,
                used: None,
                limit: None,
                unlimited: false,
                resets_at: parse_reset(&bucket["resetTime"]),
            });
        }
    }
    windows.sort_by(|a, b| a.key.cmp(&b.key));
    windows
}

// resetTime is an RFC3339 string on the live backend; fall back to epoch seconds
// (numeric or numeric string) so a format shift doesn't silently drop resets.
fn parse_reset(v: &Value) -> Option<DateTime<Utc>> {
    if let Some(s) = v.as_str() {
        if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
            return Some(dt.with_timezone(&Utc));
        }
        if let Ok(secs) = s.parse::<i64>() {
            return Utc.timestamp_opt(secs, 0).single();
        }
        return None;
    }
    v.as_i64()
        .and_then(|secs| Utc.timestamp_opt(secs, 0).single())
}

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

    const FIXTURE: &str = include_str!(concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/tests/fixtures/antigravity_quota_summary.json"
    ));

    fn fixture() -> Value {
        serde_json::from_str(FIXTURE).unwrap()
    }

    #[test]
    fn parses_quota_summary_fixture() {
        let w = parse_quota_summary(&fixture());
        // 4 buckets, but 3p-5h omits remainingFraction → 3 windows.
        assert_eq!(w.len(), 3);
        let keys: Vec<&str> = w.iter().map(|x| x.key.as_str()).collect();
        assert_eq!(keys, ["5h:gemini", "7d:3p", "7d:gemini"]);

        let g5 = w.iter().find(|x| x.key == "5h:gemini").unwrap();
        assert!((g5.used_pct - 10.0).abs() < 1e-9);
        assert_eq!(g5.kind, WindowKind::Rolling);
        assert_eq!(g5.unit, Unit::Percent);
        assert!(g5.resets_at.is_some());
        assert_eq!(g5.used, None);
        assert_eq!(g5.limit, None);
        assert!(!g5.unlimited);

        let gw = w.iter().find(|x| x.key == "7d:gemini").unwrap();
        assert_eq!(gw.used_pct, 25.0);

        let tp = w.iter().find(|x| x.key == "7d:3p").unwrap();
        assert_eq!(tp.used_pct, 0.0);

        // The fraction-less 3p-5h bucket produced no window.
        assert!(w.iter().all(|x| x.key != "5h:3p"));
    }

    #[test]
    fn skips_bucket_missing_fraction() {
        let v = serde_json::json!({ "groups": [ { "buckets": [
            { "bucketId": "a-5h", "window": "5h", "remainingFraction": 0.5, "resetTime": "2026-07-19T05:00:00Z" },
            { "bucketId": "b-5h", "window": "5h", "resetTime": "2026-07-19T05:00:00Z" }
        ] } ] });
        let w = parse_quota_summary(&v);
        assert_eq!(w.len(), 1);
        assert_eq!(w[0].key, "5h:a");
    }

    #[test]
    fn resets_at_accepts_epoch_seconds() {
        let v = serde_json::json!({ "groups": [ { "buckets": [
            { "bucketId": "x-5h", "window": "5h", "remainingFraction": 0.5, "resetTime": 1784500000 }
        ] } ] });
        let w = parse_quota_summary(&v);
        assert_eq!(w.len(), 1);
        assert!(w[0].resets_at.is_some());
    }

    #[test]
    fn plan_from_load_uses_current_not_paid_tier() {
        let v = serde_json::json!({
            "currentTier": { "name": "Antigravity", "id": "free-tier" },
            "paidTier": { "name": "Google AI Pro" }
        });
        assert_eq!(plan_from_load(&v), Some("Antigravity".to_string()));
    }
}