quotch 0.5.3

use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};

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

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

// ChatGPT's usage endpoint — verified against a live 200 on a ChatGPT Go account.
const USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";

pub struct Codex;

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

    fn discover(&self) -> Vec<Account> {
        // Emit the default account only when auth.json exists. A wholly absent
        // file means the user doesn't use Codex 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 read, no network here.
        let path = auth_path();
        if !path.exists() {
            return vec![];
        }
        vec![Account {
            provider: "codex",
            id: "default".into(),
            source: CredentialSource::FilePath(path),
            label: None,
            display: None,
        }]
    }

    fn fetch(&self, acct: &Account) -> Result<Snapshot, FetchError> {
        // Offline fallback fires ONLY on AuthMissing — the gap main.rs's stale
        // cache doesn't cover. Codex access tokens are short-lived (~1h); when one
        // expires the live call 401s and reads as "not logged in" even though the
        // user IS logged in (auth.json + refresh_token present). We deliberately
        // don't refresh (rotation risk), so we fall back to Codex's own offline
        // usage cache in the session rollout files, marked STALE. Network/Parse
        // still propagate so main.rs's stale-cache logic keeps working.
        match self.fetch_live(acct) {
            Ok(s) => Ok(s),
            Err(FetchError::AuthMissing) => offline_snapshot(acct).ok_or(FetchError::AuthMissing),
            Err(e) => Err(e),
        }
    }
}

impl Codex {
    fn fetch_live(&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 auth: Value = std::fs::read_to_string(path)
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok())
            .ok_or(FetchError::AuthMissing)?;

        let (access, account_id, id_token) = parse_auth(&auth)?;

        // The usage call needs an account id. Prefer the explicit tokens.account_id;
        // otherwise derive it from the id_token JWT. Neither → AuthMissing.
        let account_id = account_id
            .or_else(|| id_token.as_deref().and_then(account_id_from_jwt))
            .ok_or(FetchError::AuthMissing)?;

        // ponytail: same 3s connect / 8s overall split as the other providers — 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 cache.
        // No token refresh in v1: OpenAI refresh tokens rotate single-use, so
        // refreshing would mean writing auth.json back — out of scope. An expired
        // access token just 401s → AuthMissing, and the next `codex` run refreshes.
        let agent = ureq::AgentBuilder::new()
            .timeout_connect(Duration::from_secs(3))
            .timeout(Duration::from_secs(8))
            .build();

        // access/account_id ride in headers ONLY, never logged or persisted.
        let resp = match agent
            .get(USAGE_URL)
            .set("Authorization", &format!("Bearer {access}"))
            .set("ChatGPT-Account-Id", &account_id)
            .set("User-Agent", concat!("quotch/", env!("CARGO_PKG_VERSION")))
            .set("Accept", "application/json")
            .call()
        {
            Ok(r) => r,
            Err(ureq::Error::Status(code, _)) => {
                return Err(match code {
                    401 | 403 => FetchError::AuthMissing,
                    429 => FetchError::RateLimited,
                    other => FetchError::Network(format!("http {other}")),
                });
            }
            Err(ureq::Error::Transport(t)) => return Err(FetchError::Network(t.to_string())),
        };

        let body = resp
            .into_string()
            .map_err(|e| FetchError::Network(e.to_string()))?;
        let json: Value =
            serde_json::from_str(&body).map_err(|e| FetchError::Parse(e.to_string()))?;

        Ok(Snapshot {
            provider: "codex".into(),
            account: acct.id.clone(),
            label: acct.label.clone(),
            plan: plan(&json, id_token.as_deref()),
            windows: parse_windows(&json),
            fetched_at: Utc::now(),
            status: Status::Ok,
            error: None,
            raw: Some(json),
        })
    }
}

// $CODEX_HOME/auth.json when CODEX_HOME is set and non-empty, else ~/.codex/auth.json.
fn auth_path() -> PathBuf {
    match std::env::var("CODEX_HOME") {
        Ok(home) if !home.is_empty() => PathBuf::from(home).join("auth.json"),
        _ => std::env::home_dir()
            .unwrap_or_default()
            .join(".codex")
            .join("auth.json"),
    }
}

// Extracted from `fetch` so auth.json parsing is testable without a live HTTP
// call. tokens.access_token is required (missing → AuthMissing); account_id and
// id_token are optional. OPENAI_API_KEY is ignored entirely — api-key mode is
// unsupported. Empty strings are treated as absent so they never poison headers.
fn parse_auth(auth: &Value) -> Result<(String, Option<String>, Option<String>), FetchError> {
    let tokens = &auth["tokens"];
    let access = tokens["access_token"]
        .as_str()
        .filter(|s| !s.is_empty())
        .ok_or(FetchError::AuthMissing)?
        .to_string();
    let account_id = tokens["account_id"]
        .as_str()
        .filter(|s| !s.is_empty())
        .map(str::to_string);
    let id_token = tokens["id_token"]
        .as_str()
        .filter(|s| !s.is_empty())
        .map(str::to_string);
    Ok((access, account_id, id_token))
}

// plan_type from the usage response, falling back to the id_token JWT's
// chatgpt_plan_type claim (threaded through so a response without plan_type on
// some tiers still labels the plan).
fn plan(raw: &Value, id_token: Option<&str>) -> Option<String> {
    raw["plan_type"]
        .as_str()
        .map(str::to_string)
        .or_else(|| id_token.and_then(plan_from_jwt))
}

// The account id lives under the OpenAI auth namespace claim; some tokens carry
// it top-level instead. Return the first that resolves.
fn account_id_from_jwt(jwt: &str) -> Option<String> {
    let claims = jwt_payload(jwt)?;
    claims["https://api.openai.com/auth"]["chatgpt_account_id"]
        .as_str()
        .or_else(|| claims["chatgpt_account_id"].as_str())
        .map(str::to_string)
}

fn plan_from_jwt(jwt: &str) -> Option<String> {
    let claims = jwt_payload(jwt)?;
    claims["https://api.openai.com/auth"]["chatgpt_plan_type"]
        .as_str()
        .or_else(|| claims["chatgpt_plan_type"].as_str())
        .map(str::to_string)
}

// Decode a JWT's payload (segment [1]) as JSON. Signature is never verified — we
// only read our own token's claims to route the quota call.
fn jwt_payload(jwt: &str) -> Option<Value> {
    let segment = jwt.split('.').nth(1)?;
    let bytes = b64url_decode(segment)?;
    serde_json::from_slice(&bytes).ok()
}

// ponytail: base64url (no padding) decode, hand-rolled to avoid pulling the
// base64 crate into Cargo.toml just for one JWT segment. Tolerates stray '='
// padding. Any non-alphabet byte aborts to None.
fn b64url_decode(input: &str) -> Option<Vec<u8>> {
    fn sextet(c: u8) -> Option<u32> {
        Some(match c {
            b'A'..=b'Z' => u32::from(c - b'A'),
            b'a'..=b'z' => u32::from(c - b'a') + 26,
            b'0'..=b'9' => u32::from(c - b'0') + 52,
            b'-' => 62,
            b'_' => 63,
            _ => return None,
        })
    }
    let mut out = Vec::with_capacity(input.len() * 3 / 4);
    let mut acc = 0u32;
    let mut bits = 0u32;
    for &c in input.as_bytes() {
        if c == b'=' {
            continue;
        }
        acc = (acc << 6) | sextet(c)?;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((acc >> bits) as u8);
        }
    }
    Some(out)
}

// PURE, the plan-portability core. The response carries rate_limit.primary_window
// and .secondary_window, but their cadences DIFFER by plan (Go's primary is 30d;
// Plus/Pro's primary is 5h). So we derive each window's key from its DURATION and
// iterate both slots generically. Decoded defensively via Value lookups so a
// garbage secondary never discards a valid primary. Order: primary then secondary
// (natural response order — not sorted).
fn parse_windows(raw: &Value) -> Vec<Window> {
    let rate_limit = &raw["rate_limit"];
    let mut windows = Vec::new();
    for (slot, fallback_key) in [
        ("primary_window", "primary"),
        ("secondary_window", "secondary"),
    ] {
        if let Some(w) = slot_to_window(&rate_limit[slot], fallback_key) {
            windows.push(w);
        }
    }
    windows
}

fn slot_to_window(slot: &Value, fallback_key: &str) -> Option<Window> {
    // No used_percent → no usage signal; skip rather than fabricate 0%. This also
    // silently skips null/absent slots (Go's secondary_window is null).
    let used_pct = slot["used_percent"].as_f64()?;
    Some(Window {
        key: window_key(slot["limit_window_seconds"].as_i64(), fallback_key),
        kind: WindowKind::Rolling,
        unit: Unit::Percent,
        used_pct,
        used: None,
        limit: None,
        unlimited: false,
        resets_at: parse_reset(slot["reset_at"].as_i64()),
    })
}

// Key from the window duration. Known sizes map to canonical keys; unknown
// positive durations get a human-facing derived key; an absent/0 duration falls
// back to the slot name so we never silently mislabel a window we can't measure.
fn window_key(seconds: Option<i64>, fallback_key: &str) -> String {
    match seconds {
        Some(n) if n > 0 => duration_key(n),
        _ => fallback_key.to_string(),
    }
}

// The duration→key mapping, shared by the live (limit_window_seconds) and offline
// (window_minutes × 60) paths so the two can never drift. Callers guarantee n > 0.
fn duration_key(seconds: i64) -> String {
    match seconds {
        18000 => "5h".to_string(),
        604800 => "7d".to_string(),
        2592000 => "monthly".to_string(),
        n if n % 86400 == 0 => format!("{}d", n / 86400),
        n => format!("{}h", n / 3600),
    }
}

// reset_at is epoch seconds; 0/absent/negative → None.
fn parse_reset(secs: Option<i64>) -> Option<DateTime<Utc>> {
    match secs {
        Some(s) if s > 0 => DateTime::from_timestamp(s, 0),
        _ => None,
    }
}

// $CODEX_HOME/sessions when CODEX_HOME is set and non-empty, else ~/.codex/sessions.
// Mirrors auth_path()'s resolution.
fn sessions_dir() -> PathBuf {
    match std::env::var("CODEX_HOME") {
        Ok(home) if !home.is_empty() => PathBuf::from(home).join("sessions"),
        _ => std::env::home_dir()
            .unwrap_or_default()
            .join(".codex")
            .join("sessions"),
    }
}

// Recursively find the newest rollout-*.jsonl under `dir` (by mtime). std-only
// depth-first walk — no walkdir dep. Unreadable dirs/entries are skipped, not
// fatal. Returns the winning (path, mtime), or None if the tree holds none.
fn newest_rollout(dir: &Path) -> Option<(PathBuf, SystemTime)> {
    let mut best: Option<(PathBuf, SystemTime)> = None;
    for entry in std::fs::read_dir(dir).ok()?.flatten() {
        let Ok(ft) = entry.file_type() else { continue };
        if ft.is_dir() {
            if let Some(cand) = newest_rollout(&entry.path())
                && best.as_ref().is_none_or(|(_, m)| cand.1 > *m)
            {
                best = Some(cand);
            }
        } else if ft.is_file() {
            let name = entry.file_name();
            let name = name.to_string_lossy();
            if name.starts_with("rollout-")
                && name.ends_with(".jsonl")
                && let Ok(mtime) = entry.metadata().and_then(|md| md.modified())
                && best.as_ref().is_none_or(|(_, m)| mtime > *m)
            {
                best = Some((entry.path(), mtime));
            }
        }
    }
    best
}

// Pull the rate_limits object out of one rollout JSON line. Live wham/usage and
// offline rollout lines disagree on shape, so try the event-msg nesting first,
// then a top-level fallback. Non-object → None.
fn line_rate_limits(line: &str) -> Option<Value> {
    let v: Value = serde_json::from_str(line).ok()?;
    for rl in [&v["payload"]["rate_limits"], &v["rate_limits"]] {
        if rl.is_object() {
            return Some(rl.clone());
        }
    }
    None
}

// The offline fallback: reconstruct a STALE snapshot from Codex's own cached
// usage in the newest session rollout file. Any missing piece → None (caller then
// surfaces the original AuthMissing). Strictly read-only: we never write or
// refresh anything.
fn offline_snapshot(acct: &Account) -> Option<Snapshot> {
    let (path, mtime) = newest_rollout(&sessions_dir())?;

    // Keep the LAST line carrying a rate_limits object — the most recent usage
    // Codex recorded in that session.
    let contents = std::fs::read_to_string(&path).ok()?;
    let rate_limits = contents.lines().rev().find_map(line_rate_limits)?;

    let windows = parse_offline_windows(&rate_limits);

    // Age the snapshot to the rollout's mtime so staleness renders honestly;
    // fall back to now only if the timestamp is somehow unreadable.
    let fetched_at = DateTime::<Utc>::from(mtime);
    let plan = rate_limits["plan_type"].as_str().map(str::to_string);

    Some(Snapshot {
        provider: "codex".into(),
        account: acct.id.clone(),
        label: acct.label.clone(),
        plan,
        windows,
        fetched_at,
        status: Status::Stale,
        error: None,
        raw: Some(rate_limits),
    })
}

// PURE, tested. Offline rate_limits use primary/secondary (not *_window),
// window_minutes (not limit_window_seconds) and resets_at (not reset_at). Iterate
// both slots generically, deriving each key from its duration via the SAME
// duration_key() the live path uses. primary first, then secondary.
fn parse_offline_windows(rl: &Value) -> Vec<Window> {
    let mut windows = Vec::new();
    for (slot, fallback_key) in [("primary", "primary"), ("secondary", "secondary")] {
        let w = &rl[slot];
        if !w.is_object() {
            continue;
        }
        // No used_percent → no usage signal; skip rather than fabricate 0%.
        let Some(used_percent) = w["used_percent"].as_f64() else {
            continue;
        };
        let key = match w["window_minutes"].as_i64() {
            Some(m) if m > 0 => duration_key(m * 60),
            _ => fallback_key.to_string(),
        };
        windows.push(Window {
            key,
            kind: WindowKind::Rolling,
            unit: Unit::Percent,
            used_pct: used_percent.max(0.0),
            used: None,
            limit: None,
            unlimited: false,
            resets_at: parse_reset(w["resets_at"].as_i64()),
        });
    }
    windows
}

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

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

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

    // base64url (no padding) encode — test-only helper to build JWT fixtures.
    fn b64url_encode(data: &[u8]) -> String {
        const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
        let mut out = String::new();
        for chunk in data.chunks(3) {
            let b0 = u32::from(chunk[0]);
            let b1 = u32::from(*chunk.get(1).unwrap_or(&0));
            let b2 = u32::from(*chunk.get(2).unwrap_or(&0));
            let n = (b0 << 16) | (b1 << 8) | b2;
            out.push(ALPHABET[((n >> 18) & 63) as usize] as char);
            out.push(ALPHABET[((n >> 12) & 63) as usize] as char);
            if chunk.len() > 1 {
                out.push(ALPHABET[((n >> 6) & 63) as usize] as char);
            }
            if chunk.len() > 2 {
                out.push(ALPHABET[(n & 63) as usize] as char);
            }
        }
        out
    }

    fn jwt_with_payload(payload: &Value) -> String {
        format!("hdr.{}.sig", b64url_encode(payload.to_string().as_bytes()))
    }

    #[test]
    fn parses_go_fixture_single_monthly_window() {
        let w = parse_windows(&fixture());
        assert_eq!(w.len(), 1);
        assert_eq!(w[0].key, "monthly");
        assert_eq!(w[0].used_pct, 91.0);
        assert_eq!(w[0].kind, WindowKind::Rolling);
        assert_eq!(w[0].unit, Unit::Percent);
        assert!(w[0].resets_at.is_some());
        assert_eq!(w[0].used, None);
        assert_eq!(w[0].limit, None);
    }

    #[test]
    fn extracts_plan_from_fixture() {
        assert_eq!(plan(&fixture(), None), Some("go".to_string()));
    }

    // The plan-upgrade regression: the SAME code that yields one "monthly" window
    // for Go must yield 5h + 7d for Plus/Pro, in response order.
    #[test]
    fn parses_plus_pro_two_windows_in_order() {
        let v = serde_json::json!({
            "rate_limit": {
                "primary_window": {
                    "used_percent": 37,
                    "limit_window_seconds": 18000,
                    "reset_at": 1737000000
                },
                "secondary_window": {
                    "used_percent": 12,
                    "limit_window_seconds": 604800,
                    "reset_at": 1737500000
                }
            }
        });
        let w = parse_windows(&v);
        assert_eq!(w.len(), 2);
        assert_eq!(w[0].key, "5h");
        assert_eq!(w[0].used_pct, 37.0);
        assert!(w[0].resets_at.is_some());
        assert_eq!(w[1].key, "7d");
        assert_eq!(w[1].used_pct, 12.0);
        assert!(w[1].resets_at.is_some());
    }

    #[test]
    fn skips_window_missing_used_percent() {
        let v = serde_json::json!({
            "rate_limit": {
                "primary_window": {
                    "limit_window_seconds": 18000,
                    "reset_at": 1737000000
                },
                "secondary_window": null
            }
        });
        assert!(parse_windows(&v).is_empty());
    }

    #[test]
    fn derives_key_from_unknown_durations() {
        let v = serde_json::json!({
            "rate_limit": {
                "primary_window": { "used_percent": 10, "limit_window_seconds": 43200 },
                "secondary_window": { "used_percent": 20, "limit_window_seconds": 259200 }
            }
        });
        let w = parse_windows(&v);
        assert_eq!(w.len(), 2);
        assert_eq!(w[0].key, "12h");
        assert_eq!(w[1].key, "3d");
    }

    #[test]
    fn falls_back_to_slot_name_when_duration_absent() {
        let v = serde_json::json!({
            "rate_limit": {
                "primary_window": { "used_percent": 5 },
                "secondary_window": null
            }
        });
        let w = parse_windows(&v);
        assert_eq!(w.len(), 1);
        assert_eq!(w[0].key, "primary");
    }

    #[test]
    fn account_id_from_jwt_reads_namespaced_claim() {
        let jwt = jwt_with_payload(&serde_json::json!({
            "https://api.openai.com/auth": { "chatgpt_account_id": "acct-xyz" }
        }));
        assert_eq!(account_id_from_jwt(&jwt), Some("acct-xyz".to_string()));

        let missing = jwt_with_payload(&serde_json::json!({ "sub": "someone" }));
        assert_eq!(account_id_from_jwt(&missing), None);
    }

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

    // Go plan offline: one 30d/monthly primary at 91%, secondary null.
    #[test]
    fn parses_go_offline_single_monthly_window() {
        let rl: Value = serde_json::from_str(OFFLINE_FIXTURE).unwrap();
        let w = parse_offline_windows(&rl);
        assert_eq!(w.len(), 1);
        assert_eq!(w[0].key, "monthly");
        assert_eq!(w[0].used_pct, 91.0);
        assert_eq!(w[0].kind, WindowKind::Rolling);
        assert_eq!(w[0].unit, Unit::Percent);
        assert!(w[0].resets_at.is_some());
        assert_eq!(w[0].used, None);
        assert_eq!(w[0].limit, None);
    }

    // Plus/Pro offline: window_minutes drives the SAME keys as the live path
    // (300min→5h, 10080min→7d), primary then secondary.
    #[test]
    fn parses_plus_pro_offline_two_windows_in_order() {
        let rl = serde_json::json!({
            "primary": { "used_percent": 40, "window_minutes": 300, "resets_at": 1737000000 },
            "secondary": { "used_percent": 12, "window_minutes": 10080, "resets_at": 1737500000 }
        });
        let w = parse_offline_windows(&rl);
        assert_eq!(w.len(), 2);
        assert_eq!(w[0].key, "5h");
        assert_eq!(w[0].used_pct, 40.0);
        assert!(w[0].resets_at.is_some());
        assert_eq!(w[1].key, "7d");
        assert_eq!(w[1].used_pct, 12.0);
        assert!(w[1].resets_at.is_some());
    }

    #[test]
    fn skips_offline_window_missing_used_percent() {
        let rl = serde_json::json!({
            "primary": { "window_minutes": 300, "resets_at": 1737000000 },
            "secondary": null
        });
        assert!(parse_offline_windows(&rl).is_empty());
    }
}