//! Codex / "Sign in with ChatGPT" credentials.
//!
//! Lets a user with an existing ChatGPT/Codex subscription run Lucid Train on
//! that subscription instead of a pay-per-token API key. Two ways in:
//!
//!  1. Import an existing Codex CLI login from `~/.codex/auth.json`.
//!  2. Run the OAuth 2.0 PKCE "Sign in with ChatGPT" flow ourselves (browser +
//!     a one-shot `localhost:1455` callback).
//!
//! Both yield a short-lived `access_token` (a JWT) plus a `refresh_token`. The
//! token is sent as a bearer to the Codex subscription endpoint
//! (`https://chatgpt.com/backend-api/codex/responses`) together with the
//! `chatgpt-account-id` header. Tokens are refreshed against
//! `https://auth.openai.com/oauth/token` and written back to the SAME
//! `auth.json` — refresh tokens are single-use, so sharing one file with the
//! Codex CLI keeps both clients valid instead of invalidating each other.
//!
//! NOTE: the subscription endpoint is undocumented and only accepts callers
//! presenting Codex's own OAuth client id + `originator: codex_cli_rs`. This is
//! reverse-engineered interop, subject to OpenAI's terms; treat it as such.

use anyhow::{anyhow, bail, Context, Result};
use base64::Engine;
use serde::Deserialize;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::io::Write;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};

/// Codex's public OAuth client id (the same one the official CLI uses — the
/// subscription endpoint only accepts tokens minted for this client).
pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
const AUTH_BASE: &str = "https://auth.openai.com";
const REDIRECT_PORT: u16 = 1455;
const SCOPE: &str = "openid profile email offline_access";

/// Credentials for the Codex subscription endpoint.
#[derive(Debug, Clone, Default)]
pub struct CodexTokens {
    pub access_token: String,
    pub refresh_token: String,
    pub id_token: String,
    /// ChatGPT account id, sent as the `chatgpt-account-id` header.
    pub account_id: String,
}

/// Shape of `~/.codex/auth.json` (the inner token bag).
#[derive(Debug, Deserialize, Default)]
struct AuthFileTokens {
    #[serde(default)]
    access_token: String,
    #[serde(default)]
    refresh_token: String,
    #[serde(default)]
    id_token: String,
    #[serde(default)]
    account_id: String,
}

#[derive(Debug, Deserialize, Default)]
struct AuthFile {
    #[serde(default)]
    tokens: AuthFileTokens,
    // Some shapes store the token bag flat; accept that too.
    #[serde(default)]
    access_token: String,
    #[serde(default)]
    refresh_token: String,
    #[serde(default)]
    id_token: String,
    #[serde(default)]
    account_id: String,
}

/// `$CODEX_HOME/auth.json`, else `~/.codex/auth.json`.
pub fn auth_path() -> PathBuf {
    if let Ok(home) = std::env::var("CODEX_HOME") {
        if !home.is_empty() {
            return PathBuf::from(home).join("auth.json");
        }
    }
    dirs::home_dir().unwrap_or_default().join(".codex").join("auth.json")
}

/// Load credentials from an existing Codex login, if present and usable.
pub fn load() -> Option<CodexTokens> {
    let raw = std::fs::read_to_string(auth_path()).ok()?;
    let parsed: AuthFile = serde_json::from_str(&raw).ok()?;
    // Prefer the nested `tokens` bag; fall back to flat fields.
    let (access, refresh, id, mut account) = if !parsed.tokens.access_token.is_empty() {
        (
            parsed.tokens.access_token,
            parsed.tokens.refresh_token,
            parsed.tokens.id_token,
            parsed.tokens.account_id,
        )
    } else {
        (parsed.access_token, parsed.refresh_token, parsed.id_token, parsed.account_id)
    };
    if access.is_empty() {
        return None;
    }
    if account.is_empty() {
        account = account_id_from_jwt(&access)
            .or_else(|| account_id_from_jwt(&id))
            .unwrap_or_default();
    }
    Some(CodexTokens { access_token: access, refresh_token: refresh, id_token: id, account_id: account })
}

/// Persist refreshed tokens back to `auth.json`, preserving any other keys and
/// keeping owner-only (0600) permissions. Stays compatible with the Codex CLI.
pub fn save(tokens: &CodexTokens) -> Result<()> {
    let path = auth_path();
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).ok();
    }
    // Merge into the existing file so we don't clobber fields we don't model.
    let mut root: Value = std::fs::read_to_string(&path)
        .ok()
        .and_then(|s| serde_json::from_str(&s).ok())
        .unwrap_or_else(|| json!({}));
    if !root.is_object() {
        root = json!({});
    }
    root["tokens"] = json!({
        "access_token": tokens.access_token,
        "refresh_token": tokens.refresh_token,
        "id_token": tokens.id_token,
        "account_id": tokens.account_id,
    });
    root["last_refresh"] = json!(rfc3339_now());
    let body = serde_json::to_string_pretty(&root)?;

    let mut opts = std::fs::OpenOptions::new();
    opts.write(true).create(true).truncate(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        opts.mode(0o600);
    }
    let mut f = opts.open(&path).with_context(|| format!("writing {}", path.display()))?;
    f.write_all(body.as_bytes())?;
    Ok(())
}

/// True if the access token is missing or its `exp` is within 60s.
pub fn is_expired(tokens: &CodexTokens) -> bool {
    if tokens.access_token.is_empty() {
        return true;
    }
    match jwt_exp(&tokens.access_token) {
        Some(exp) => now_secs() + 60 >= exp,
        None => false, // opaque token — assume valid, refresh reactively on 401
    }
}

/// Exchange the refresh token for a fresh access token and persist it.
pub async fn refresh(tokens: &mut CodexTokens, http: &reqwest::Client) -> Result<()> {
    if tokens.refresh_token.is_empty() {
        bail!("no refresh token — run /login codex again");
    }
    let resp = http
        .post(format!("{AUTH_BASE}/oauth/token"))
        .json(&json!({
            "grant_type": "refresh_token",
            "client_id": CLIENT_ID,
            "refresh_token": tokens.refresh_token,
            "scope": SCOPE,
        }))
        .send()
        .await
        .context("refresh request failed")?;
    let status = resp.status();
    let v: Value = resp.json().await.context("refresh response not JSON")?;
    if !status.is_success() {
        bail!("token refresh HTTP {status}: {v}");
    }
    apply_token_response(tokens, &v);
    save(tokens).ok();
    Ok(())
}

/// Re-read `auth.json` and adopt fresher tokens if the file carries a non-empty
/// refresh token. The Codex CLI shares this file and may have refreshed (and
/// rotated the single-use refresh token) since we loaded ours; adopting the
/// on-disk copy keeps us from refreshing with an already-spent token.
pub fn reload(tokens: &mut CodexTokens) {
    if let Some(disk) = load() {
        if !disk.refresh_token.is_empty() {
            *tokens = disk;
        }
    }
}

/// Recover a usable Codex login from disk: load it, then refresh when the token
/// is expired (or when `force`, e.g. after a 401 the access token is
/// known-dead). Returns live tokens (also written back to disk). Errors when
/// there's no login on disk or the refresh fails — the caller then runs the
/// browser sign-in. This is what keeps a Codex session alive across restarts
/// instead of falling back to an API-key prompt.
pub async fn ensure_valid(http: &reqwest::Client, force: bool) -> Result<CodexTokens> {
    let mut tokens = load().ok_or_else(|| anyhow!("no Codex login found on disk"))?;
    if force || is_expired(&tokens) {
        refresh(&mut tokens, http).await?;
    }
    Ok(tokens)
}

/// Run the browser PKCE login and return fresh tokens (also written to disk).
/// `on_url` is invoked once with the authorize URL (before the browser opens) so
/// callers can show it as a fallback if `open_browser` silently fails.
pub async fn login(http: &reqwest::Client, on_url: impl FnOnce(&str)) -> Result<CodexTokens> {
    let (verifier, challenge) = pkce_pair()?;
    let state = rand_token(24)?;
    let redirect = format!("http://localhost:{REDIRECT_PORT}/auth/callback");

    let auth_url = {
        let mut u = url::Url::parse(&format!("{AUTH_BASE}/oauth/authorize"))?;
        u.query_pairs_mut()
            .append_pair("response_type", "code")
            .append_pair("client_id", CLIENT_ID)
            .append_pair("redirect_uri", &redirect)
            .append_pair("scope", SCOPE)
            .append_pair("code_challenge", &challenge)
            .append_pair("code_challenge_method", "S256")
            .append_pair("id_token_add_organizations", "true")
            .append_pair("state", &state);
        u.to_string()
    };

    // Start the callback listener BEFORE opening the browser to avoid a race.
    let listener = tokio::net::TcpListener::bind(("127.0.0.1", REDIRECT_PORT))
        .await
        .with_context(|| format!("cannot bind localhost:{REDIRECT_PORT} (is another login running?)"))?;

    on_url(&auth_url);
    open_browser(&auth_url);

    // Wait (with a timeout) for OpenAI to redirect back with ?code=...&state=...
    let code = tokio::time::timeout(std::time::Duration::from_secs(300), wait_for_code(&listener, &state))
        .await
        .map_err(|_| anyhow!("login timed out after 5 minutes"))??;

    // Exchange the authorization code for tokens.
    let resp = http
        .post(format!("{AUTH_BASE}/oauth/token"))
        .json(&json!({
            "grant_type": "authorization_code",
            "client_id": CLIENT_ID,
            "code": code,
            "redirect_uri": redirect,
            "code_verifier": verifier,
        }))
        .send()
        .await
        .context("token exchange failed")?;
    let status = resp.status();
    let v: Value = resp.json().await.context("token response not JSON")?;
    if !status.is_success() {
        bail!("token exchange HTTP {status}: {v}");
    }
    let mut tokens = CodexTokens::default();
    apply_token_response(&mut tokens, &v);
    if tokens.access_token.is_empty() {
        bail!("login succeeded but no access_token was returned");
    }
    save(&tokens).ok();
    Ok(tokens)
}

/// A random UUIDv4-shaped session id for the `session_id` request header.
pub fn session_id() -> String {
    let mut b = [0u8; 16];
    if getrandom::getrandom(&mut b).is_err() {
        // Fall back to time-derived bytes; the header just needs to be present.
        let t = now_secs().to_le_bytes();
        b[..8].copy_from_slice(&t);
    }
    b[6] = (b[6] & 0x0f) | 0x40; // version 4
    b[8] = (b[8] & 0x3f) | 0x80; // variant
    let h = |r: &[u8]| r.iter().map(|x| format!("{x:02x}")).collect::<String>();
    format!("{}-{}-{}-{}-{}", h(&b[0..4]), h(&b[4..6]), h(&b[6..8]), h(&b[8..10]), h(&b[10..16]))
}

/// A human-readable label for the signed-in account (for `/status`).
pub fn account_label(tokens: &CodexTokens) -> String {
    if tokens.account_id.is_empty() {
        "ChatGPT account".to_string()
    } else {
        tokens.account_id.clone()
    }
}

// ---------- internals ----------

fn apply_token_response(tokens: &mut CodexTokens, v: &Value) {
    if let Some(s) = v["access_token"].as_str() {
        if !s.is_empty() {
            tokens.access_token = s.to_string();
        }
    }
    // Refresh tokens rotate; keep the old one only if none came back.
    if let Some(s) = v["refresh_token"].as_str() {
        if !s.is_empty() {
            tokens.refresh_token = s.to_string();
        }
    }
    if let Some(s) = v["id_token"].as_str() {
        if !s.is_empty() {
            tokens.id_token = s.to_string();
        }
    }
    if tokens.account_id.is_empty() {
        tokens.account_id = account_id_from_jwt(&tokens.access_token)
            .or_else(|| account_id_from_jwt(&tokens.id_token))
            .unwrap_or_default();
    }
}

/// Wait for the OAuth redirect, validate `state`, return the `code`. Replies to
/// the browser with a small success/error page so the tab is self-explanatory.
async fn wait_for_code(listener: &tokio::net::TcpListener, expected_state: &str) -> Result<String> {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    loop {
        let (mut sock, _) = listener.accept().await.context("callback accept failed")?;
        let mut buf = vec![0u8; 8192];
        let n = sock.read(&mut buf).await.unwrap_or(0);
        let req = String::from_utf8_lossy(&buf[..n]);
        // First line: "GET /auth/callback?code=...&state=... HTTP/1.1"
        let path = req.lines().next().and_then(|l| l.split_whitespace().nth(1)).unwrap_or("");
        if !path.starts_with("/auth/callback") {
            // Ignore favicon/etc. probes; keep listening.
            let _ = sock.write_all(http_response("Lucid Train", "You can close this tab.")).await;
            continue;
        }
        let full = format!("http://localhost{path}");
        let parsed = url::Url::parse(&full).ok();
        let mut code = None;
        let mut state = None;
        let mut err = None;
        if let Some(u) = &parsed {
            for (k, val) in u.query_pairs() {
                match k.as_ref() {
                    "code" => code = Some(val.into_owned()),
                    "state" => state = Some(val.into_owned()),
                    "error" => err = Some(val.into_owned()),
                    _ => {}
                }
            }
        }
        if let Some(e) = err {
            let _ = sock.write_all(http_response("Login failed", &e)).await;
            bail!("authorization denied: {e}");
        }
        if state.as_deref() != Some(expected_state) {
            let _ = sock.write_all(http_response("Login failed", "state mismatch")).await;
            bail!("state mismatch (possible CSRF) — login aborted");
        }
        match code {
            Some(c) if !c.is_empty() => {
                let _ = sock
                    .write_all(http_response("Signed in to Lucid Train", "Return to your terminal."))
                    .await;
                return Ok(c);
            }
            _ => {
                let _ = sock.write_all(http_response("Login failed", "no authorization code")).await;
                bail!("no authorization code in callback");
            }
        }
    }
}

fn http_response(title: &str, msg: &str) -> &'static [u8] {
    // We can't easily build a borrowed static with dynamic content, so leak a
    // tiny formatted page. Login happens at most a handful of times per process.
    let body = format!(
        "<!doctype html><html><head><meta charset=utf-8><title>{title}</title>\
         <style>body{{background:#1a1b26;color:#c0caf5;font:16px/1.6 ui-monospace,Menlo,monospace;\
         display:flex;align-items:center;justify-content:center;height:100vh;margin:0}}\
         .c{{text-align:center}}b{{color:#7aa2f7}}</style></head>\
         <body><div class=c><h2><b>&gt;_</b> {title}</h2><p>{msg}</p></div></body></html>"
    );
    let resp =
        format!("HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body.len(), body);
    Box::leak(resp.into_bytes().into_boxed_slice())
}

fn open_browser(url: &str) {
    #[cfg(target_os = "macos")]
    let cmd = "open";
    #[cfg(target_os = "windows")]
    let cmd = "explorer";
    #[cfg(all(unix, not(target_os = "macos")))]
    let cmd = "xdg-open";
    let _ = std::process::Command::new(cmd)
        .arg(url)
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .spawn();
}

/// PKCE: random verifier + S256 challenge (base64url, no padding).
fn pkce_pair() -> Result<(String, String)> {
    let verifier = rand_token(64)?;
    let digest = Sha256::digest(verifier.as_bytes());
    let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
    Ok((verifier, challenge))
}

/// Cryptographically random URL-safe token of roughly `bytes` of entropy.
fn rand_token(bytes: usize) -> Result<String> {
    let mut raw = vec![0u8; bytes];
    getrandom::getrandom(&mut raw).map_err(|e| anyhow!("rng failure: {e}"))?;
    Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw))
}

/// Decode a JWT payload (middle segment) into JSON claims.
fn jwt_claims(token: &str) -> Option<Value> {
    let payload = token.split('.').nth(1)?;
    let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(payload)
        .or_else(|_| base64::engine::general_purpose::STANDARD_NO_PAD.decode(payload))
        .ok()?;
    serde_json::from_slice(&bytes).ok()
}

fn jwt_exp(token: &str) -> Option<u64> {
    jwt_claims(token)?["exp"].as_u64()
}

/// Codex stores the account id under the `https://api.openai.com/auth` claim.
fn account_id_from_jwt(token: &str) -> Option<String> {
    let claims = jwt_claims(token)?;
    let auth = &claims["https://api.openai.com/auth"];
    for key in ["chatgpt_account_id", "account_id", "organization_id"] {
        if let Some(s) = auth[key].as_str() {
            if !s.is_empty() {
                return Some(s.to_string());
            }
        }
    }
    None
}

fn now_secs() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)
}

/// Format the current time as RFC3339 UTC (no chrono dependency).
fn rfc3339_now() -> String {
    let secs = now_secs() as i64;
    // Days since epoch and seconds within the day.
    let days = secs.div_euclid(86_400);
    let rem = secs.rem_euclid(86_400);
    let (hh, mm, ss) = (rem / 3600, (rem % 3600) / 60, rem % 60);
    // Civil date from days (Howard Hinnant's algorithm).
    let z = days + 719_468;
    let era = z.div_euclid(146_097);
    let doe = z - era * 146_097;
    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    format!("{y:04}-{m:02}-{d:02}T{hh:02}:{mm:02}:{ss:02}Z")
}

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

    /// Serialise tests that mutate the process-wide `CODEX_HOME` env var.
    static ENV_LOCK: Mutex<()> = Mutex::new(());

    fn b64url(v: &Value) -> String {
        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(serde_json::to_vec(v).unwrap())
    }

    fn fake_jwt(claims: Value) -> String {
        format!("{}.{}.{}", b64url(&json!({"alg":"none"})), b64url(&claims), "sig")
    }

    #[test]
    fn pkce_challenge_is_s256_of_verifier() {
        let (verifier, challenge) = pkce_pair().unwrap();
        let expect = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
        assert_eq!(challenge, expect);
        assert!(!verifier.contains('=') && !verifier.contains('+') && !verifier.contains('/'));
    }

    #[test]
    fn extracts_exp_and_account_id_from_jwt() {
        let tok = fake_jwt(json!({
            "exp": 4_000_000_000u64,
            "https://api.openai.com/auth": { "chatgpt_account_id": "acct_123" }
        }));
        assert_eq!(jwt_exp(&tok), Some(4_000_000_000));
        assert_eq!(account_id_from_jwt(&tok).as_deref(), Some("acct_123"));
    }

    #[test]
    fn is_expired_respects_exp() {
        let fresh = fake_jwt(json!({ "exp": now_secs() + 3600 }));
        let stale = fake_jwt(json!({ "exp": now_secs().saturating_sub(10) }));
        assert!(!is_expired(&CodexTokens { access_token: fresh, ..Default::default() }));
        assert!(is_expired(&CodexTokens { access_token: stale, ..Default::default() }));
        assert!(is_expired(&CodexTokens::default()));
    }

    #[test]
    fn parses_nested_and_flat_auth_json() {
        let nested = r#"{"tokens":{"access_token":"a","refresh_token":"r","account_id":"acct_x"}}"#;
        let parsed: AuthFile = serde_json::from_str(nested).unwrap();
        assert_eq!(parsed.tokens.access_token, "a");
        assert_eq!(parsed.tokens.account_id, "acct_x");

        let flat = r#"{"access_token":"a2","refresh_token":"r2"}"#;
        let parsed: AuthFile = serde_json::from_str(flat).unwrap();
        assert_eq!(parsed.access_token, "a2");
    }

    #[test]
    fn reload_adopts_fresher_on_disk_tokens() {
        let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // The Codex CLI shares auth.json and may rotate the refresh token; a
        // reload must adopt the on-disk copy so we don't refresh with a spent
        // token. A file with no refresh token is ignored (keeps ours).
        let dir = std::env::temp_dir().join(format!("lucid-codex-reload-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let prev = std::env::var("CODEX_HOME").ok();
        std::env::set_var("CODEX_HOME", &dir);

        std::fs::write(
            dir.join("auth.json"),
            r#"{"tokens":{"access_token":"disk-access","refresh_token":"disk-refresh","account_id":"acct_disk"}}"#,
        )
        .unwrap();
        let mut tokens = CodexTokens {
            access_token: "stale".into(),
            refresh_token: "stale-refresh".into(),
            ..Default::default()
        };
        reload(&mut tokens);
        assert_eq!(tokens.access_token, "disk-access");
        assert_eq!(tokens.refresh_token, "disk-refresh");
        assert_eq!(tokens.account_id, "acct_disk");

        // A file without a refresh token must not clobber our live one.
        std::fs::write(dir.join("auth.json"), r#"{"tokens":{"access_token":"x"}}"#).unwrap();
        let mut keep = CodexTokens { refresh_token: "keep-me".into(), ..Default::default() };
        reload(&mut keep);
        assert_eq!(keep.refresh_token, "keep-me");

        if let Some(p) = prev {
            std::env::set_var("CODEX_HOME", p);
        } else {
            std::env::remove_var("CODEX_HOME");
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn ensure_valid_errs_without_a_login() {
        let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // With no auth.json the caller must get an error (and then fall back to
        // the browser sign-in) rather than a bogus token.
        let dir = std::env::temp_dir().join(format!("lucid-codex-empty-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let prev = std::env::var("CODEX_HOME").ok();
        std::env::set_var("CODEX_HOME", &dir);

        let http = reqwest::Client::new();
        assert!(ensure_valid(&http, false).await.is_err());

        if let Some(p) = prev {
            std::env::set_var("CODEX_HOME", p);
        } else {
            std::env::remove_var("CODEX_HOME");
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn rfc3339_now_has_expected_shape() {
        let s = rfc3339_now();
        assert_eq!(s.len(), 20, "got {s}");
        assert!(s.ends_with('Z') && s.contains('T'));
    }
}
