quotch 0.5.1

use std::path::PathBuf;
use std::time::Duration;

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> {
        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(18000) => "5h".to_string(),
        Some(604800) => "7d".to_string(),
        Some(2592000) => "monthly".to_string(),
        Some(n) if n > 0 => {
            if n % 86400 == 0 {
                format!("{}d", n / 86400)
            } else {
                format!("{}h", n / 3600)
            }
        }
        _ => fallback_key.to_string(),
    }
}

// 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,
    }
}

#[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);
    }
}