magi-code 0.96.1

use super::readiness::{
    self as auth, AuthProviderRecord, AuthState, CredentialReadiness, ProviderCredential,
};
use crate::config::McPaths;
use crate::http_body::{DEFAULT_BOUNDED_BODY_MAX_BYTES, read_bounded_response_text};
use crate::providers::OPENAI_CODEX_PROVIDER;
use anyhow::Context;
use chrono::Utc;
use serde::Deserialize;
use std::time::Duration;

pub(crate) const OPENAI_CODEX_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
pub(crate) const OPENAI_CODEX_RELOGIN_GUIDANCE: &str =
    "openai-codex OAuth credentials need re-login; run /login openai-codex";
pub(crate) const OPENAI_CODEX_REDIRECT_URI: &str = "http://localhost:1455/auth/callback";
const TOKEN_URL: &str = "https://auth.openai.com/oauth/token";

#[derive(Debug, Clone)]
pub(crate) struct NormalizedToken {
    pub(crate) access: String,
    pub(crate) refresh: Option<String>,
    pub(crate) expires: Option<i64>,
    pub(crate) account_id: String,
}

#[derive(Deserialize)]
struct TokenResponse {
    access_token: Option<String>,
    refresh_token: Option<String>,
    expires_in: Option<i64>,
    #[serde(rename = "accountId")]
    account_id_camel: Option<String>,
    account_id: Option<String>,
    id_token: Option<String>,
}

pub(crate) fn codex_credential_from_store(paths: &McPaths) -> anyhow::Result<ProviderCredential> {
    codex_credential_from_store_with_exchange(paths, refresh_codex_token)
}

pub(crate) fn codex_credential_from_store_with_exchange(
    paths: &McPaths,
    exchange: impl FnOnce(&str) -> anyhow::Result<NormalizedToken>,
) -> anyhow::Result<ProviderCredential> {
    let stored = auth::read_auth_store(paths)?;
    let expected_generation = stored.provider_generation(OPENAI_CODEX_PROVIDER);
    let Some(source_record) = stored.auth().providers.get(OPENAI_CODEX_PROVIDER).cloned() else {
        anyhow::bail!("missing OAuth auth for provider 'openai-codex'; run /login openai-codex")
    };
    let AuthProviderRecord::OAuth {
        access,
        refresh,
        expires,
        account_id,
    } = &source_record
    else {
        anyhow::bail!("provider 'openai-codex' requires OAuth auth; run /login openai-codex")
    };
    let now = Utc::now().timestamp();
    match auth::classify_codex_oauth_record(
        access,
        refresh.as_deref(),
        *expires,
        account_id.as_deref(),
        now,
    ) {
        CredentialReadiness::Ready => {
            return Ok(ProviderCredential::OAuth {
                access: access.clone(),
                account_id: account_id.clone(),
            });
        }
        CredentialReadiness::Refreshable => {}
        CredentialReadiness::Missing | CredentialReadiness::Invalid => {
            anyhow::bail!(OPENAI_CODEX_RELOGIN_GUIDANCE)
        }
    }
    let refresh = refresh
        .as_deref()
        .filter(|value| !value.is_empty())
        .ok_or_else(|| anyhow::anyhow!(OPENAI_CODEX_RELOGIN_GUIDANCE))?;
    let token = match exchange(refresh) {
        Ok(token) => token,
        Err(error) => {
            return current_ready_credential_after_failed_exchange(
                paths,
                expected_generation,
                error,
            );
        }
    };
    complete_codex_refresh(paths, expected_generation, source_record, token)
}

pub(crate) fn force_refresh_codex_credential_from_store(
    paths: &McPaths,
) -> anyhow::Result<ProviderCredential> {
    force_refresh_codex_credential_from_store_with_exchange(paths, refresh_codex_token)
}

fn force_refresh_codex_credential_from_store_with_exchange(
    paths: &McPaths,
    exchange: impl FnOnce(&str) -> anyhow::Result<NormalizedToken>,
) -> anyhow::Result<ProviderCredential> {
    let stored = auth::read_auth_store(paths)?;
    let expected_generation = stored.provider_generation(OPENAI_CODEX_PROVIDER);
    let Some(source_record) = stored.auth().providers.get(OPENAI_CODEX_PROVIDER).cloned() else {
        anyhow::bail!(OPENAI_CODEX_RELOGIN_GUIDANCE)
    };
    let AuthProviderRecord::OAuth { refresh, .. } = &source_record else {
        anyhow::bail!(OPENAI_CODEX_RELOGIN_GUIDANCE)
    };
    let refresh = refresh
        .as_deref()
        .filter(|value| !value.is_empty())
        .ok_or_else(|| anyhow::anyhow!(OPENAI_CODEX_RELOGIN_GUIDANCE))?;
    let token = match exchange(refresh) {
        Ok(token) => token,
        Err(error) => {
            return current_ready_credential_after_failed_exchange(
                paths,
                expected_generation,
                error,
            );
        }
    };
    complete_codex_refresh(paths, expected_generation, source_record, token)
}

fn complete_codex_refresh(
    paths: &McPaths,
    expected_generation: u64,
    source_record: AuthProviderRecord,
    token: NormalizedToken,
) -> anyhow::Result<ProviderCredential> {
    let previous_refresh = match &source_record {
        AuthProviderRecord::OAuth { refresh, .. } => refresh.clone(),
        AuthProviderRecord::ApiKey { .. } => None,
    };
    let replacement = oauth_record_from_token(token.clone(), previous_refresh);
    match auth::replace_provider_auth_if_matches(
        paths,
        OPENAI_CODEX_PROVIDER,
        expected_generation,
        &source_record,
        replacement,
    )? {
        auth::ConditionalAuthUpdate::Applied => Ok(provider_credential_from_token(token)),
        auth::ConditionalAuthUpdate::Current(current_store) => {
            codex_credential_from_current_store(&current_store)
        }
    }
}

fn current_ready_credential_after_failed_exchange(
    paths: &McPaths,
    expected_generation: u64,
    exchange_error: anyhow::Error,
) -> anyhow::Result<ProviderCredential> {
    let current_store = match auth::read_auth_store(paths) {
        Ok(store) => store,
        Err(_) => return Err(exchange_error),
    };
    if current_store.provider_generation(OPENAI_CODEX_PROVIDER) <= expected_generation {
        return Err(exchange_error);
    }
    match codex_credential_from_current_store(&current_store) {
        Ok(credential) => Ok(credential),
        Err(_) => Err(exchange_error),
    }
}

fn codex_credential_from_current_store(
    current_store: &auth::AuthStore,
) -> anyhow::Result<ProviderCredential> {
    let Some(record) = current_store.auth().providers.get(OPENAI_CODEX_PROVIDER) else {
        anyhow::bail!("missing OAuth auth for provider 'openai-codex'; run /login openai-codex")
    };
    let AuthProviderRecord::OAuth {
        access,
        refresh,
        expires,
        account_id,
    } = record
    else {
        anyhow::bail!("provider 'openai-codex' requires OAuth auth; run /login openai-codex")
    };
    // A concurrent refresh may have committed a new Ready record. Never treat a current
    // expired/Refreshable record as usable just because it has an access token and account id.
    let readiness = auth::classify_codex_oauth_record(
        access,
        refresh.as_deref(),
        *expires,
        account_id.as_deref(),
        Utc::now().timestamp(),
    );
    if readiness == CredentialReadiness::Ready {
        return Ok(ProviderCredential::OAuth {
            access: access.clone(),
            account_id: account_id.clone(),
        });
    }
    anyhow::bail!(
        "Codex OAuth credentials changed during refresh and are not currently usable; {OPENAI_CODEX_RELOGIN_GUIDANCE}"
    )
}

fn provider_credential_from_token(token: NormalizedToken) -> ProviderCredential {
    ProviderCredential::OAuth {
        access: token.access,
        account_id: Some(token.account_id),
    }
}

pub(super) fn oauth_record_from_token(
    token: NormalizedToken,
    previous_refresh: Option<String>,
) -> AuthProviderRecord {
    AuthProviderRecord::OAuth {
        access: token.access,
        refresh: token.refresh.or(previous_refresh),
        expires: token.expires,
        account_id: Some(token.account_id),
    }
}

pub(crate) fn refreshed_codex_auth_state(
    paths: &McPaths,
    current: &AuthState,
) -> anyhow::Result<AuthState> {
    if current.provider() != OPENAI_CODEX_PROVIDER {
        return Ok(current.clone());
    }
    let credential = codex_credential_from_store(paths)?;
    Ok(AuthState::Ready {
        provider: OPENAI_CODEX_PROVIDER.to_string(),
        credential,
    })
}

pub(crate) fn persist_codex_token(paths: &McPaths, token: NormalizedToken) -> anyhow::Result<()> {
    auth::update_auth(paths, OPENAI_CODEX_PROVIDER, |auth| {
        let previous_refresh = match auth.providers.get(OPENAI_CODEX_PROVIDER) {
            Some(AuthProviderRecord::OAuth { refresh, .. }) => refresh.clone(),
            _ => None,
        };
        auth.providers.insert(
            OPENAI_CODEX_PROVIDER.to_string(),
            oauth_record_from_token(token, previous_refresh),
        );
    })?;
    Ok(())
}

pub(crate) fn exchange_codex_code(verifier: &str, code: &str) -> anyhow::Result<NormalizedToken> {
    let body = [
        ("grant_type", "authorization_code"),
        ("client_id", OPENAI_CODEX_CLIENT_ID),
        ("redirect_uri", OPENAI_CODEX_REDIRECT_URI),
        ("code", code),
        ("code_verifier", verifier),
    ];
    post_token_form(&body)
}

pub(crate) fn refresh_codex_token(refresh: &str) -> anyhow::Result<NormalizedToken> {
    let body = [
        ("grant_type", "refresh_token"),
        ("client_id", OPENAI_CODEX_CLIENT_ID),
        ("refresh_token", refresh),
    ];
    post_token_form(&body)
}

fn post_token_form(body: &[(&str, &str)]) -> anyhow::Result<NormalizedToken> {
    let response = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(30))
        .build()?
        .post(TOKEN_URL)
        .form(body)
        .send()?;
    let status = response.status();
    if !status.is_success() {
        anyhow::bail!("OpenAI Codex OAuth token exchange failed with status {status}")
    }
    let text = read_bounded_response_text(response, DEFAULT_BOUNDED_BODY_MAX_BYTES)
        .with_context(|| "OpenAI Codex OAuth token response body read failed")?;
    let response = serde_json::from_str::<TokenResponse>(&text)?;
    normalize_token_response(response)
}

fn normalize_token_response(response: TokenResponse) -> anyhow::Result<NormalizedToken> {
    let access = response
        .access_token
        .filter(|value| !value.is_empty())
        .ok_or_else(|| anyhow::anyhow!("OAuth token response missing access token"))?;
    let account_id = response
        .account_id_camel
        .or(response.account_id)
        .or_else(|| auth::extract_oauth_account_id_from_jwt(&access))
        .or_else(|| {
            response
                .id_token
                .as_deref()
                .and_then(auth::extract_oauth_account_id_from_jwt)
        })
        .filter(|value| !value.is_empty())
        .ok_or_else(|| anyhow::anyhow!("OAuth token response missing ChatGPT account id"))?;
    let expires = response
        .expires_in
        .map(|seconds| {
            Utc::now()
                .timestamp()
                .checked_add(seconds)
                .ok_or_else(|| anyhow::anyhow!("OAuth token response expires_in is too large"))
        })
        .transpose()?;
    Ok(NormalizedToken {
        access,
        refresh: response.refresh_token.filter(|value| !value.is_empty()),
        expires,
        account_id,
    })
}