//! OAuth sessions for servers and the CLI: the atproto OAuth client
//! flow (PAR + PKCE + DPoP), a session saved per user DID in the state dir (the sealed setting
//! "oauth:<did>"), refreshed by host::Session (one refresh per process, and across processes under the
//! store's "oauth:<did>" lease) and revoked on sign-out.
//!
//! Two clients, as knotless has them (crates/knotless/src/oauth.rs): the loopback client
//! (`http://localhost?redirect_uri=...&scope=...`, public: `atlibs login`, `atlibs serve` with no
//! --origin) and a confidential client on a public origin (`<origin>/oauth/client-metadata.json`,
//! private_key_jwt with the state dir's key, `<origin>/oauth/jwks.json`). One form POST serves both:
//! the confidential one adds its client assertion. The pieces are atproto-crates': PKCE
//! (atproto_oauth::pkce), state, DPoP proofs (atproto_oauth::dpop), client assertions
//! (atproto_oauth::jwt), keys (atproto_identity::key) and the authorization server's metadata checks
//! (atproto_oauth::resources::validate_authorization_server). Its workflow functions are
//! confidential-only and bound to their own reqwest client, so the requests go through host::Fetch
//! (tests and the conformance bridge run without a network; public_url keeps DID-named hosts off the
//! host's network).

use std::collections::HashMap;
use std::sync::Mutex;

use atproto_identity::key::{KeyData, KeyType, generate_key, identify_key};
use atproto_oauth::jwt::{Claims, Header, JoseClaims, mint};
use atproto_oauth::resources::{AuthorizationServer, validate_authorization_server};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde_json::{Map, Value, json};

use crate::host::{Fetch, Response, public_url};
use crate::store::{Store, random_bytes};
use crate::{js_string, now_ms};

/// The sealed setting an OAuth session of `did` is saved in, and the lease its refreshes take.
pub fn setting(did: &str) -> String {
    format!("oauth:{did}")
}

/// A client of authorization servers: loopback (public), or confidential with its key.
#[derive(Clone)]
pub struct Client {
    pub id: String,
    pub redirect_uri: String,
    pub scope: String,
    pub key: Option<KeyData>,
}

impl Client {
    /// The loopback client: its redirect and scope are in its id. The id's redirect leaves out the
    /// port (any port the listener gets matches); `redirect_uri` is the one this login listens on.
    pub fn loopback(redirect_uri: &str, scope: &str) -> Self {
        let mut named = url::Url::parse(redirect_uri).expect("a loopback redirect URL");
        let _ = named.set_port(None);
        let query = url::form_urlencoded::Serializer::new(String::new()).append_pair("redirect_uri", named.as_str()).append_pair("scope", scope).finish();
        Self { id: format!("http://localhost?{query}"), redirect_uri: redirect_uri.into(), scope: scope.into(), key: None }
    }

    /// The confidential client of a public origin: `<origin><path>client-metadata.json`, redirecting
    /// to `<origin><path>callback`.
    pub fn confidential(origin: &str, path: &str, scope: &str, key: KeyData) -> Self {
        let base = format!("{}{path}", origin.trim_end_matches('/'));
        Self { id: format!("{base}client-metadata.json"), redirect_uri: format!("{base}callback"), scope: scope.into(), key: Some(key) }
    }

    /// Its public keys (atproto_oauth::jwk), what `<path>jwks.json` serves.
    pub fn jwks(&self) -> Value {
        let keys: Vec<Value> = self.key.iter().filter_map(|k| atproto_oauth::jwk::generate(k).ok()).filter_map(|k| serde_json::to_value(k).ok()).collect();
        json!({ "keys": keys })
    }

    /// The client metadata document served at `id` (confidential clients; a loopback id describes itself).
    pub fn metadata(&self, name: &str) -> Value {
        let origin = url::Url::parse(&self.id).map(|u| u.origin().ascii_serialization()).unwrap_or_default();
        json!({
            "client_id": self.id,
            "client_name": name,
            "client_uri": origin,
            "redirect_uris": [self.redirect_uri],
            "scope": self.scope,
            "grant_types": ["authorization_code", "refresh_token"],
            "response_types": ["code"],
            "application_type": "web",
            "token_endpoint_auth_method": "private_key_jwt",
            "token_endpoint_auth_signing_alg": "ES256",
            "jwks_uri": self.id.replace("client-metadata.json", "jwks.json"),
            "dpop_bound_access_tokens": true,
        })
    }

    /// The form fields naming this client to `issuer`: a confidential one's carry an assertion signed by its key.
    fn fields(&self, issuer: &str) -> Result<Vec<(String, String)>, String> {
        let mut fields = vec![("client_id".to_string(), self.id.clone())];
        if let Some(key) = &self.key {
            let now = now_ms() / 1000;
            let claims = Claims::new(JoseClaims {
                issuer: Some(self.id.clone()),
                subject: Some(self.id.clone()),
                audience: Some(issuer.into()),
                json_web_token_id: Some(random()),
                issued_at: Some(now),
                expiration: Some(now + 60),
                ..Default::default()
            });
            let header = Header::try_from(key.clone()).map_err(|e| e.to_string())?;
            fields.push(("client_assertion_type".into(), "urn:ietf:params:oauth:client-assertion-type:jwt-bearer".into()));
            fields.push(("client_assertion".into(), mint(key, &header, &claims).map_err(|e| e.to_string())?));
        }
        Ok(fields)
    }
}

/// 16 random bytes, base64url.
fn random() -> String {
    URL_SAFE_NO_PAD.encode(random_bytes::<16>())
}

/// The confidential client's key: a P-256 key sealed in the state dir ("oauthClientKey"), made once.
pub async fn client_key(store: &Store) -> Result<KeyData, String> {
    let _lease = store.lock("oauthClientKey", std::time::Duration::from_secs(10)).await?;
    if let Some(saved) = store.setting("oauthClientKey")?.as_ref().and_then(Value::as_str) {
        return identify_key(saved).map_err(|e| e.to_string());
    }
    let key = generate_key(KeyType::P256Private).map_err(|e| e.to_string())?;
    store.set_setting("oauthClientKey", &json!(key.to_string()))?;
    Ok(key)
}

/// The DPoP nonce each server last sent, by origin (one process's; a stale one costs a retry).
static NONCES: Mutex<Option<HashMap<String, String>>> = Mutex::new(None);

fn origin(url: &str) -> String {
    url::Url::parse(url).map(|u| u.origin().ascii_serialization()).unwrap_or_default()
}

fn nonce(url: &str) -> Option<String> {
    NONCES.lock().unwrap().as_ref()?.get(&origin(url)).cloned()
}

/// Keeps the DPoP-Nonce a server's answer carries.
pub fn keep_nonce(url: &str, response: &Response) {
    if let Some((_, n)) = response.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case("dpop-nonce")) {
        NONCES.lock().unwrap().get_or_insert_with(HashMap::new).insert(origin(url), n.clone());
    }
}

/// A DPoP proof that `key` holds the request `method url` (with `token`'s ath on a resource request),
/// carrying the server's last nonce.
pub fn proof(key: &KeyData, method: &str, url: &str, token: Option<&str>) -> Result<String, String> {
    use atproto_oauth::dpop::{auth_dpop, dpop_with_nonce, request_dpop};
    let made = match (nonce(url), token) {
        (Some(n), token) => dpop_with_nonce(key, method, url, token, &n),
        (None, Some(t)) => request_dpop(key, method, url, t),
        (None, None) => auth_dpop(key, method, url),
    };
    made.map(|(jwt, ..)| jwt).map_err(|e| e.to_string())
}

/// Whether a resource server refused a request for want of the current nonce (send it again).
pub fn nonce_refused(response: &Response) -> bool {
    response.status == 401 && header(response, "www-authenticate").is_some_and(|v| v.contains("use_dpop_nonce"))
}

/// Whether a resource server refused the access token itself (refresh it).
pub fn token_refused(response: &Response) -> bool {
    response.status == 401 && header(response, "www-authenticate").is_some_and(|v| v.contains("invalid_token"))
}

fn header<'a>(response: &'a Response, name: &str) -> Option<&'a str> {
    response.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str())
}

fn allow_private() -> bool {
    std::env::var("ATLIBS_ALLOW_PRIVATE").is_ok_and(|v| v == "1")
}

async fn get_json(fetch: &dyn Fetch, url: &str) -> Result<Value, String> {
    let response = fetch.fetch("GET", &public_url(url, allow_private())?, vec![], None).await?;
    if !(200..300).contains(&response.status) {
        return Err(format!("{url}: {}", response.status));
    }
    serde_json::from_slice(&response.body).map_err(|e| format!("{url}: {e}"))
}

/// The authorization server of accounts on `pds` (its metadata, checked by atproto-oauth).
async fn server(fetch: &dyn Fetch, pds: &str) -> Result<Value, String> {
    let pds = pds.trim_end_matches('/');
    let resource = get_json(fetch, &format!("{pds}/.well-known/oauth-protected-resource")).await.map_err(|e| format!("{pds} does not say where to sign in with OAuth: {e}"))?;
    let issuer = resource["authorization_servers"][0].as_str().ok_or_else(|| format!("{pds} names no authorization server"))?.to_string();
    let doc = get_json(fetch, &format!("{issuer}/.well-known/oauth-authorization-server")).await?;
    let checked: AuthorizationServer = serde_json::from_value(doc.clone()).map_err(|e| format!("{issuer}: {e}"))?;
    validate_authorization_server(&checked, &issuer).map_err(|e| format!("{issuer}: {e}"))?;
    Ok(doc)
}

/// POSTs `form` from `client` to `url` at `issuer`, with a DPoP proof of `key`; a use_dpop_nonce
/// answer is retried once with the nonce it brought.
async fn post(fetch: &dyn Fetch, url: &str, key: &KeyData, client: &Client, issuer: &str, form: &[(&str, &str)]) -> Result<Value, String> {
    let url = public_url(url, allow_private())?;
    for retried in [false, true] {
        let body = {
            let mut body = url::form_urlencoded::Serializer::new(String::new());
            body.extend_pairs(form.iter().copied());
            body.extend_pairs(client.fields(issuer)?);
            body.finish()
        };
        let headers = vec![("content-type".into(), "application/x-www-form-urlencoded".into()), ("dpop".into(), proof(key, "POST", &url, None)?)];
        let response = fetch.fetch("POST", &url, headers, Some(body.into_bytes())).await?;
        keep_nonce(&url, &response);
        let value: Value = serde_json::from_slice(&response.body).unwrap_or(Value::Null);
        if (200..300).contains(&response.status) {
            return Ok(value);
        }
        if !retried && value["error"] == "use_dpop_nonce" {
            continue;
        }
        return Err(format!("{}: {}", value["error"].as_str().unwrap_or("refused"), value["error_description"].as_str().unwrap_or_default()));
    }
    unreachable!("the second try returns")
}

/// Starts signing in as `did` (whose PDS is `pds`; `hint` is what the user typed): pushes the
/// authorization request. Returns the page to send the browser to and the pending request (sealed
/// by whoever keeps it: it holds the DPoP key and PKCE verifier) for `complete`.
pub async fn begin(fetch: &dyn Fetch, client: &Client, did: &str, pds: &str, hint: &str) -> Result<(String, Value), String> {
    let server = server(fetch, pds).await?;
    let issuer = js_string(server.get("issuer"));
    let key = generate_key(KeyType::P256Private).map_err(|e| e.to_string())?;
    let (verifier, challenge) = atproto_oauth::pkce::generate();
    let state = atproto_oauth::workflow::generate_state();
    let pushed = post(
        fetch,
        &js_string(server.get("pushed_authorization_request_endpoint")),
        &key,
        client,
        &issuer,
        &[
            ("response_type", "code"),
            ("code_challenge", &challenge),
            ("code_challenge_method", "S256"),
            ("redirect_uri", &client.redirect_uri),
            ("scope", &client.scope),
            ("state", &state),
            ("login_hint", hint.trim_start_matches('@')),
        ],
    )
    .await?;
    let request_uri = pushed["request_uri"].as_str().ok_or("the authorization server answered no request_uri")?;
    let mut url = url::Url::parse(&js_string(server.get("authorization_endpoint"))).map_err(|e| e.to_string())?;
    url.query_pairs_mut().append_pair("client_id", &client.id).append_pair("request_uri", request_uri);
    let pending = json!({
        "state": state, "did": did, "pds": pds, "issuer": issuer, "verifier": verifier, "dpopKey": key.to_string(),
        "tokenEndpoint": server["token_endpoint"], "revocationEndpoint": server["revocation_endpoint"],
        "redirectUri": client.redirect_uri, "scope": client.scope, "createdAt": now_ms(),
    });
    Ok((url.to_string(), pending))
}

/// Ends a sign-in: the callback's query (`params`) against its pending request. The browser that
/// came back must be the one that left (the caller checks that: a cookie, or a one-shot loopback
/// listener). Returns the session to save (Session::save_oauth).
pub async fn complete(fetch: &dyn Fetch, client: &Client, pending: &Value, params: &Map<String, Value>) -> Result<Value, String> {
    let issuer = js_string(pending.get("issuer"));
    let param = |k: &str| params.get(k).and_then(Value::as_str);
    if param("state") != pending["state"].as_str() || param("iss").is_some_and(|i| i != issuer) {
        return Err("the browser came back from a sign-in this one did not start".into());
    }
    let Some(code) = param("code") else {
        return Err(format!("the sign-in was refused: {}", param("error_description").or(param("error")).unwrap_or("no reason given")));
    };
    if now_ms().saturating_sub(pending["createdAt"].as_u64().unwrap_or(0)) > 10 * 60 * 1000 {
        return Err("the sign-in took too long; start again".into());
    }
    let key = identify_key(pending["dpopKey"].as_str().unwrap_or_default()).map_err(|e| e.to_string())?;
    let form = [("grant_type", "authorization_code"), ("code", code), ("redirect_uri", pending["redirectUri"].as_str().unwrap_or_default()), ("code_verifier", pending["verifier"].as_str().unwrap_or_default())];
    let tokens = post(fetch, &js_string(pending.get("tokenEndpoint")), &key, client, &issuer, &form).await?;
    let did = js_string(pending.get("did"));
    if tokens["sub"].as_str() != Some(did.as_str()) {
        return Err(format!("signed in as {}, not {did}", js_string(tokens.get("sub"))));
    }
    session_of(&tokens, pending, client)
}

/// The saved session from a token response: the shape Session keeps ("oauth:<did>").
fn session_of(tokens: &Value, from: &Value, client: &Client) -> Result<Value, String> {
    if !tokens["token_type"].as_str().is_some_and(|t| t.eq_ignore_ascii_case("dpop")) {
        return Err("the authorization server did not bind the tokens to DPoP".into());
    }
    let (Some(access), Some(refresh)) = (tokens["access_token"].as_str(), tokens["refresh_token"].as_str()) else {
        return Err("the token response has no access or refresh token".into());
    };
    Ok(json!({
        "did": from["did"], "pds": from["pds"], "issuer": from["issuer"], "dpopKey": from["dpopKey"],
        "tokenEndpoint": from["tokenEndpoint"], "revocationEndpoint": from["revocationEndpoint"],
        "clientId": client.id, "authMethod": if client.key.is_some() { "private_key_jwt" } else { "none" },
        "scope": tokens.get("scope").cloned().unwrap_or_else(|| from["scope"].clone()),
        "accessJwt": access, "refreshJwt": refresh,
        "expiresAt": now_ms() + tokens["expires_in"].as_u64().unwrap_or(300) * 1000,
    }))
}

/// The client a saved session was granted to (its id; a confidential one signs with the state dir's key).
pub async fn client_of(saved: &Value, store: Option<&Store>) -> Result<Client, String> {
    let key = match (saved["authMethod"].as_str(), store) {
        (Some("private_key_jwt"), Some(store)) => Some(client_key(store).await?),
        (Some("private_key_jwt"), None) => return Err("a confidential client's session needs its state dir".into()),
        _ => None,
    };
    Ok(Client { id: js_string(saved.get("clientId")), redirect_uri: String::new(), scope: js_string(saved.get("scope")), key })
}

/// New tokens for a saved session (refresh tokens are used once: Session runs this one at a time).
pub async fn refresh(fetch: &dyn Fetch, saved: &Value, store: Option<&Store>) -> Result<Value, String> {
    let client = client_of(saved, store).await?;
    let key = identify_key(saved["dpopKey"].as_str().unwrap_or_default()).map_err(|e| e.to_string())?;
    let form = [("grant_type", "refresh_token"), ("refresh_token", saved["refreshJwt"].as_str().unwrap_or_default())];
    let tokens = post(fetch, &js_string(saved.get("tokenEndpoint")), &key, &client, &js_string(saved.get("issuer")), &form).await?;
    if tokens["sub"].as_str().is_some_and(|s| Some(s) != saved["did"].as_str()) {
        return Err("the refreshed tokens are another account's".into());
    }
    session_of(&tokens, saved, &client)
}

/// Revokes a saved session's refresh token (and so its access), as far as its server lets it.
pub async fn revoke(fetch: &dyn Fetch, saved: &Value, store: Option<&Store>) -> Result<(), String> {
    let Some(endpoint) = saved["revocationEndpoint"].as_str() else { return Ok(()) };
    let client = client_of(saved, store).await?;
    let key = identify_key(saved["dpopKey"].as_str().unwrap_or_default()).map_err(|e| e.to_string())?;
    post(fetch, endpoint, &key, &client, &js_string(saved.get("issuer")), &[("token", saved["refreshJwt"].as_str().unwrap_or_default())]).await.map(drop)
}

/// Signs `did` out of this state dir: revokes its tokens and forgets the session.
pub async fn sign_out(fetch: &dyn Fetch, store: &Store, did: &str) -> Result<bool, String> {
    let Some(saved) = store.setting(&setting(did))? else { return Ok(false) };
    let revoked = revoke(fetch, &saved, Some(store)).await;
    store.delete(&format!("{}:settings", store.id), &setting(did))?;
    revoked.map(|_| true)
}

/// Opaque session handles (a cookie or a bearer) -> the DID they sign in as. Only their sha256 is kept.
pub fn handle_key(handle: &str) -> String {
    format!("oauth-handle:{}", crate::wasm::sha256_hex(handle.as_bytes()))
}

pub fn new_handle(store: &Store, did: &str) -> Result<String, String> {
    let handle = URL_SAFE_NO_PAD.encode(random_bytes::<32>());
    store.put(&format!("{}:kv", store.id), &handle_key(&handle), &json!({ "did": did, "createdAt": now_ms() }))?;
    Ok(handle)
}

pub fn did_of_handle(store: &Store, handle: &str) -> Result<Option<String>, String> {
    Ok(store.get(&format!("{}:kv", store.id), &handle_key(handle))?.and_then(|v| v["did"].as_str().map(String::from)))
}

pub fn drop_handle(store: &Store, handle: &str) -> Result<(), String> {
    store.delete(&format!("{}:kv", store.id), &handle_key(handle)).map(drop)
}