//! SQLite-backed OAuth persistence + confidential-client bootstrap for Jacquard 0.12.
//!
//! [`SqliteAuthStore`] implements [`ClientAuthStore`] and
//! [`SessionSelector`] over the `oauth_sessions` and `oauth_auth_requests`
//! tables (see `migrations/0001_initial_schema.sql`). Each row stores a
//! structured, queryable projection of the Jacquard session / auth-request types
//! alongside a canonical JSON `payload`. Structured columns power key
//! enumeration, by-DID/any selection, and deletes without the payload; only
//! `get_session` / `get_auth_req_info` deserialize the payload.
//! DPoP key material and nonces live inside the payload — the Jacquard
//! `SessionRegistry` mutex is in-process concurrency only and is not persisted,
//! so the two tables are the *complete* durable OAuth surface.
//!
//! [`bootstrap_oauth`] constructs the confidential [`OAuthClient`], the
//! single-origin [`OAuthWebConfig`], and the private cookie-jar [`Key`] from a
//! swappable [`SecretSource`] (keyring → env → ephemeral-with-warning).
//!
//! [`SessionSelector`]: jacquard::common::session::SessionSelector

use jacquard::common::bos::BosStr;
use jacquard::common::session::{SessionHint, SessionKey, SessionSelector, SessionStoreError};
use jacquard::common::types::did::Did;
use jacquard::oauth::authstore::{ClientAuthStore, OAuthSessionMatch};
use jacquard::oauth::session::{AuthRequestData, ClientSessionData};
use sqlx::SqlitePool;

/// SQLite-backed implementation of [`ClientAuthStore`].
#[derive(Debug, Clone)]
pub struct SqliteAuthStore {
    pool: SqlitePool,
}

impl SqliteAuthStore {
    /// Wrap the shared projection pool as an OAuth auth store.
    pub fn new(pool: SqlitePool) -> Self {
        Self { pool }
    }

    /// Exact-key lookup → session match (or `None`).
    async fn match_for_key(
        &self,
        key: SessionKey,
    ) -> Result<Option<OAuthSessionMatch>, SessionStoreError> {
        Ok(self
            .get_session(&key.did, key.session_id.as_str())
            .await?
            .map(|session| OAuthSessionMatch { key, session }))
    }

    /// First session belonging to `did`, enumerated via [`ClientAuthStore::list_session_keys`].
    async fn match_for_did<D: BosStr + Send + Sync>(
        &self,
        did: &Did<D>,
    ) -> Result<Option<OAuthSessionMatch>, SessionStoreError> {
        for key in self.list_session_keys().await? {
            if key.did.as_str() == did.as_ref()
                && let Some(matched) = self.match_for_key(key).await?
            {
                return Ok(Some(matched));
            }
        }
        Ok(None)
    }
}

impl ClientAuthStore for SqliteAuthStore {
    async fn get_session<D: BosStr + Send + Sync>(
        &self,
        did: &Did<D>,
        session_id: &str,
    ) -> Result<Option<ClientSessionData>, SessionStoreError> {
        let row = sqlx::query!(
            "SELECT payload FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
            did.as_str(),
            session_id
        )
        .fetch_optional(&self.pool)
        .await
        .map_err(store_err)?;
        match row {
            Some(r) => Ok(Some(serde_json::from_str(&r.payload)?)),
            None => Ok(None),
        }
    }

    async fn upsert_session(&self, session: ClientSessionData) -> Result<(), SessionStoreError> {
        let now = chrono::Utc::now()
            .timestamp_nanos_opt()
            .expect("current timestamp must fit in i64 nanoseconds");
        let payload = serde_json::to_string(&session)?;
        sqlx::query!(
            "INSERT INTO oauth_sessions
               (account_did, session_id, host_url, authserver_url, token_endpoint,
                revocation_endpoint, created_at, updated_at, payload)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
             ON CONFLICT(account_did, session_id) DO UPDATE SET
               host_url = excluded.host_url,
               authserver_url = excluded.authserver_url,
               token_endpoint = excluded.token_endpoint,
               revocation_endpoint = excluded.revocation_endpoint,
               updated_at = excluded.updated_at,
               payload = excluded.payload",
            session.account_did.as_str(),
            session.session_id.as_str(),
            session.host_url.as_str(),
            session.authserver_url.as_str(),
            session.authserver_token_endpoint.as_str(),
            session
                .authserver_revocation_endpoint
                .as_ref()
                .map(|s| s.as_str()),
            now,
            now,
            payload,
        )
        .execute(&self.pool)
        .await
        .map_err(store_err)?;
        Ok(())
    }

    async fn delete_session<D: BosStr + Send + Sync>(
        &self,
        did: &Did<D>,
        session_id: &str,
    ) -> Result<(), SessionStoreError> {
        sqlx::query!(
            "DELETE FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
            did.as_str(),
            session_id
        )
        .execute(&self.pool)
        .await
        .map_err(store_err)?;
        Ok(())
    }

    async fn get_auth_req_info(
        &self,
        state: &str,
    ) -> Result<Option<AuthRequestData>, SessionStoreError> {
        let row = sqlx::query!(
            "SELECT payload FROM oauth_auth_requests WHERE state = ?",
            state
        )
        .fetch_optional(&self.pool)
        .await
        .map_err(store_err)?;
        match row {
            Some(r) => Ok(Some(serde_json::from_str(&r.payload)?)),
            None => Ok(None),
        }
    }

    async fn save_auth_req_info(
        &self,
        auth_req_info: &AuthRequestData,
    ) -> Result<(), SessionStoreError> {
        let now = chrono::Utc::now()
            .timestamp_nanos_opt()
            .expect("current timestamp must fit in i64 nanoseconds");
        let payload = serde_json::to_string(auth_req_info)?;
        sqlx::query!(
            "INSERT INTO oauth_auth_requests
               (state, authserver_url, account_did, request_uri, token_endpoint,
                revocation_endpoint, created_at, payload)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?)
             ON CONFLICT(state) DO UPDATE SET
               authserver_url = excluded.authserver_url,
               account_did = excluded.account_did,
               request_uri = excluded.request_uri,
               token_endpoint = excluded.token_endpoint,
               revocation_endpoint = excluded.revocation_endpoint,
               payload = excluded.payload",
            auth_req_info.state.as_str(),
            auth_req_info.authserver_url.as_str(),
            auth_req_info.account_did.as_ref().map(|d| d.as_str()),
            auth_req_info.request_uri.as_str(),
            auth_req_info.authserver_token_endpoint.as_str(),
            auth_req_info
                .authserver_revocation_endpoint
                .as_ref()
                .map(|s| s.as_str()),
            now,
            payload,
        )
        .execute(&self.pool)
        .await
        .map_err(store_err)?;
        Ok(())
    }

    async fn delete_auth_req_info(&self, state: &str) -> Result<(), SessionStoreError> {
        sqlx::query!("DELETE FROM oauth_auth_requests WHERE state = ?", state)
            .execute(&self.pool)
            .await
            .map_err(store_err)?;
        Ok(())
    }

    async fn list_session_keys(&self) -> Result<Vec<SessionKey>, SessionStoreError> {
        let rows = sqlx::query!("SELECT account_did, session_id FROM oauth_sessions")
            .fetch_all(&self.pool)
            .await
            .map_err(store_err)?;
        let mut keys: Vec<SessionKey> = Vec::with_capacity(rows.len());
        for row in rows {
            let did = Did::new_owned(&row.account_did).map_err(store_err)?;
            keys.push(SessionKey::new(did, row.session_id));
        }
        keys.sort();
        Ok(keys)
    }
}

impl SessionSelector<OAuthSessionMatch> for SqliteAuthStore {
    type Error = SessionStoreError;

    async fn select_session<Str: BosStr + Send + Sync>(
        &self,
        hint: &SessionHint<Str>,
    ) -> Result<Option<OAuthSessionMatch>, Self::Error> {
        match hint {
            SessionHint::Any => {
                let Some(key) = self.list_session_keys().await?.into_iter().next() else {
                    return Ok(None);
                };
                self.match_for_key(key).await
            }
            SessionHint::Did(did) => self.match_for_did(did).await,
            SessionHint::Handle(_) | SessionHint::Identifier(_) => Ok(None),
            SessionHint::Key(key) => self.match_for_key(key.clone()).await,
        }
    }
}

/// Box any backend error into [`SessionStoreError::Other`].
fn store_err<E>(e: E) -> SessionStoreError
where
    E: std::error::Error + Send + Sync + 'static,
{
    SessionStoreError::Other(Box::new(e))
}

// ===========================================================================
// Confidential OAuth client bootstrap
// ===========================================================================

use anyhow::Context as _;
use axum_extra::extract::cookie::Key;
use base64::Engine as _;
use jacquard::common::deps::fluent_uri::Uri;
use jacquard::identity::PublicResolver;
use jacquard::oauth::{
    atproto::AtprotoClientMetadata, client::OAuthClient, keyset::Keyset, scopes::Scopes,
    session::ClientData,
};
use jacquard_axum::oauth::OAuthWebConfig;

/// The OAuth client, web config, and cookie-jar key built once at startup.
///
/// Constructed by [`bootstrap_oauth`] and stored in `AppState`.
pub struct OAuthBootstrap {
    /// Confidential OAuth client over the SQLite auth store + public resolver.
    pub client: OAuthClient<PublicResolver, SqliteAuthStore>,
    /// Single-origin cookie/route config (`OAuthWebConfig::default()`).
    pub config: OAuthWebConfig,
    /// Signing key for the private `jacquard_oauth_session` cookie jar.
    pub cookie_key: Key,
}

/// Build the OAuth client, web config, and cookie key.
///
/// `base_url` is the server's public origin (`POLYMODEL_BASE_URL`); the OAuth
/// `client_id` is `{base_url}/oauth-client-metadata.json` and the redirect URI
/// is `{base_url}/oauth/callback`, mirroring the jacquard-axum example. Only the
/// secret store (→ [`SqliteAuthStore`]) and key provisioning (→
/// [`SecretSource`]) differ from the example.
pub fn bootstrap_oauth(pool: SqlitePool, base_url: Option<&str>) -> anyhow::Result<OAuthBootstrap> {
    let (keyset, cookie_key) = provision_oauth_secrets()?;
    let metadata = hosted_client_metadata(base_url)?;
    let client_data = ClientData::new(keyset, metadata);
    let client: OAuthClient<PublicResolver, SqliteAuthStore> = OAuthClient::new(
        SqliteAuthStore::new(pool),
        client_data,
        reqwest::Client::new(),
    );
    Ok(OAuthBootstrap {
        client,
        config: OAuthWebConfig::default(),
        cookie_key,
    })
}

/// Build hosted (non-loopback) client metadata from the public origin.
///
/// Defaults to the deployment origin `https://polymodel.space` when
/// `POLYMODEL_BASE_URL` is unset; set that for local or alternate origins. This is
/// never the loopback client spec (`new_localhost`); it uses the hosted
/// `AtprotoClientMetadata::new` constructor for production origins.
fn hosted_client_metadata(
    base_url: Option<&str>,
) -> anyhow::Result<AtprotoClientMetadata<jacquard::common::deps::smol_str::SmolStr>> {
    let origin = base_url
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(|s| s.trim_end_matches('/').to_owned());
    let origin = match origin {
        Some(origin) => origin,
        None => {
            tracing::warn!(
                "POLYMODEL_BASE_URL is not set; defaulting the OAuth client origin to the \
                 deployment origin https://polymodel.space (used for client_id and redirect_uri). \
                 Set POLYMODEL_BASE_URL for local or alternate origins."
            );
            "https://polymodel.space".to_owned()
        }
    };
    let client_id =
        Uri::parse(format!("{origin}/oauth-client-metadata.json")).map_err(|(_e, _)| {
            anyhow::anyhow!("invalid OAuth client_id URL from POLYMODEL_BASE_URL")
        })?;
    let redirect_uri = Uri::parse(format!("{origin}/oauth/callback"))
        .map_err(|(_e, _)| anyhow::anyhow!("invalid OAuth redirect URI from POLYMODEL_BASE_URL"))?;
    let scopes = Scopes::builder()
        .atproto()
        .transition_generic()
        .build()
        .context("build OAuth scopes")?;
    Ok(AtprotoClientMetadata::new(
        vec![redirect_uri],
        client_id,
        Some(scopes),
    ))
}

/// Provision the OAuth signing keyset and cookie-jar key.
///
/// Keyset priority: (1) OS keyring — preferred, persists across restarts; on
/// first run a fresh keyset is generated and stored so subsequent restarts
/// reuse it; (2) environment override (`POLYMODEL_OAUTH_KEYSET_B64`) —
/// read-only; (3) `None` — fall back to public OAuth (no client assertion).
/// Ephemeral keysets are never generated because they change on restart and
/// poison the PDS client-assertion cache.
///
/// Cookie key: same keyring/env priority, then ephemeral generation (sessions
/// are lost on restart but the PDS is unaffected).
fn provision_oauth_secrets() -> anyhow::Result<(Option<Keyset>, Key)> {
    let keyring = KeyringSource;
    let keyset = provision_keyset(&keyring)?;
    let (cookie, generated_cookie) = provision_cookie_key(&keyring)?;
    if generated_cookie {
        tracing::warn!(
            "No durable cookie key found (keyring/env). Generated an ephemeral one — browser \
             sessions will not survive restart."
        );
    }
    Ok((keyset, cookie))
}

fn provision_keyset(keyring: &KeyringSource) -> anyhow::Result<Option<Keyset>> {
    if !cfg!(test)
        && let Some(keyset) = keyring.load_keyset()?
    {
        return Ok(Some(keyset));
    }
    if !cfg!(test)
        && let Some(keyset) = EnvSource.load_keyset()?
    {
        return Ok(Some(keyset));
    }
    // Try to generate and persist. Only return the keyset if the keyring
    // actually stored it — ephemeral keys change on restart and poison the PDS
    // client-assertion cache, so we fall back to public OAuth instead.
    if !cfg!(test) {
        let keyset = Keyset::generate_es256("polymodel").context("generate OAuth keyset")?;
        let _ = keyring.store_keyset(&keyset);
        if keyring.load_keyset()?.is_some() {
            return Ok(Some(keyset));
        }
    }
    tracing::warn!(
        "No durable OAuth keyset found (keyring/env). Using public OAuth (no client assertion). \
         Set POLYMODEL_OAUTH_KEYSET_B64 or install an OS keyring backend for DPoP client auth."
    );
    Ok(None)
}

fn provision_cookie_key(keyring: &KeyringSource) -> anyhow::Result<(Key, bool)> {
    if !cfg!(test)
        && let Some(cookie) = keyring.load_cookie_key()?
    {
        return Ok((cookie, false));
    }
    if !cfg!(test)
        && let Some(cookie) = EnvSource.load_cookie_key()?
    {
        return Ok((cookie, false));
    }

    let cookie = Key::generate();
    // Best-effort: persist to keyring so a later restart (on a host with a usable
    // secret store) reuses this instead of regenerating. The cookie key must stay
    // stable for OAuth session restore.
    if !cfg!(test) {
        let _ = keyring.store_cookie_key(&cookie);
    }
    Ok((cookie, true))
}

/// Swappable source of the OAuth signing keyset and cookie-jar key.
///
/// Sources are tried in priority order by [`provision_oauth_secrets`]; the
/// ephemeral fallback is intentionally not a `SecretSource` because it is
/// non-persistent by design.
trait SecretSource {
    fn load_keyset(&self) -> anyhow::Result<Option<Keyset>>;
    fn store_keyset(&self, keyset: &Keyset) -> anyhow::Result<()>;
    fn load_cookie_key(&self) -> anyhow::Result<Option<Key>>;
    fn store_cookie_key(&self, key: &Key) -> anyhow::Result<()>;
}

/// OS keyring-backed source. On a host without an accessible secret store, every
/// operation degrades to `None` / no-op so provisioning falls through to env.
struct KeyringSource;

#[cfg(feature = "server")]
impl KeyringSource {
    const SERVICE: &'static str = "polymodel";
    const KEYSET_USER: &'static str = "oauth-client-keyset";
    const COOKIE_USER: &'static str = "oauth-cookie-key";

    fn entry(user: &'static str) -> Option<keyring::Entry> {
        keyring::Entry::new(Self::SERVICE, user).ok()
    }
}

#[cfg(feature = "server")]
impl SecretSource for KeyringSource {
    fn load_keyset(&self) -> anyhow::Result<Option<Keyset>> {
        let Some(entry) = Self::entry(Self::KEYSET_USER) else {
            return Ok(None);
        };
        match entry.get_password() {
            Ok(b64) => {
                let bytes = decode_b64(&b64).context("keyring OAuth keyset")?;
                Ok(Some(serde_json::from_slice(&bytes).context(
                    "keyring OAuth keyset is not a valid serialized Keyset",
                )?))
            }
            // Not stored yet, or no usable secret-store backend → fall through.
            Err(_) => Ok(None),
        }
    }

    fn store_keyset(&self, keyset: &Keyset) -> anyhow::Result<()> {
        if let Some(entry) = Self::entry(Self::KEYSET_USER) {
            let bytes = serde_json::to_vec(keyset)?;
            let _ = entry.set_password(&b64::STANDARD.encode(&bytes));
        }
        Ok(())
    }

    fn load_cookie_key(&self) -> anyhow::Result<Option<Key>> {
        let Some(entry) = Self::entry(Self::COOKIE_USER) else {
            return Ok(None);
        };
        match entry.get_password() {
            Ok(b64) => decode_cookie_key(&b64).map(Some),
            Err(_) => Ok(None),
        }
    }

    fn store_cookie_key(&self, key: &Key) -> anyhow::Result<()> {
        if let Some(entry) = Self::entry(Self::COOKIE_USER) {
            let _ = entry.set_password(&b64::STANDARD.encode(key.master()));
        }
        Ok(())
    }
}
#[cfg(not(feature = "server"))]
impl SecretSource for KeyringSource {
    fn load_keyset(&self) -> anyhow::Result<Option<Keyset>> {
        Ok(None)
    }
    fn store_keyset(&self, _keyset: &Keyset) -> anyhow::Result<()> {
        Ok(())
    }
    fn load_cookie_key(&self) -> anyhow::Result<Option<Key>> {
        Ok(None)
    }
    fn store_cookie_key(&self, _key: &Key) -> anyhow::Result<()> {
        Ok(())
    }
}

/// Environment-variable source (`POLYMODEL_OAUTH_KEYSET_B64`,
/// `POLYMODEL_COOKIE_KEY_B64`). Base64 of a JSON-serialized `Keyset` and of 64
/// raw cookie-key bytes respectively. Read-only.
struct EnvSource;

impl SecretSource for EnvSource {
    fn load_keyset(&self) -> anyhow::Result<Option<Keyset>> {
        let Some(b64) = env_get("POLYMODEL_OAUTH_KEYSET_B64")? else {
            return Ok(None);
        };
        let bytes = decode_b64(&b64).context("POLYMODEL_OAUTH_KEYSET_B64")?;
        Ok(Some(serde_json::from_slice(&bytes).context(
            "POLYMODEL_OAUTH_KEYSET_B64 is not a valid serialized Keyset",
        )?))
    }

    fn store_keyset(&self, _keyset: &Keyset) -> anyhow::Result<()> {
        Ok(())
    }

    fn load_cookie_key(&self) -> anyhow::Result<Option<Key>> {
        let Some(b64) = env_get("POLYMODEL_COOKIE_KEY_B64")? else {
            return Ok(None);
        };
        decode_cookie_key(&b64)
            .context("POLYMODEL_COOKIE_KEY_B64")
            .map(Some)
    }

    fn store_cookie_key(&self, _key: &Key) -> anyhow::Result<()> {
        Ok(())
    }
}

fn env_get(name: &str) -> anyhow::Result<Option<String>> {
    match std::env::var(name) {
        Ok(v) => Ok(Some(v)),
        Err(std::env::VarError::NotPresent) => Ok(None),
        Err(std::env::VarError::NotUnicode(_)) => {
            anyhow::bail!("{name} is not valid UTF-8")
        }
    }
}

fn decode_b64(b64: &str) -> anyhow::Result<Vec<u8>> {
    Ok(b64::STANDARD.decode(b64)?)
}

fn decode_cookie_key(b64: &str) -> anyhow::Result<Key> {
    let bytes = decode_b64(b64)?;
    let arr: [u8; 64] = bytes
        .as_slice()
        .try_into()
        .map_err(|_| anyhow::anyhow!("cookie key must decode to exactly 64 bytes"))?;
    Ok(Key::from(&arr))
}

/// Short alias for standard base64 to keep the keyring/env code compact.
mod b64 {
    pub(crate) use base64::engine::general_purpose::STANDARD;
}

#[cfg(test)]
mod tests {
    //! Store behavior, not serde: keyed upsert/get, overwrite-on-conflict,
    //! delete exactness, key enumeration + by-DID querying (structured column,
    //! no payload load), `select_session` for `Key`/`Did`/`Any`/no-match,
    //! auth-request lifecycle, and cross-session isolation.

    use super::*;
    use jacquard::common::deps::fluent_uri::Uri;
    use jacquard::common::deps::smol_str::SmolStr;
    use jacquard::oauth::keyset::Keyset;
    use jacquard::oauth::scopes::Scopes;
    use jacquard::oauth::session::{DpopClientData, DpopReqData};
    use jacquard::oauth::types::{OAuthTokenType, TokenSet};
    use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
    use std::str::FromStr;

    const ALICE: &str = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa";
    const BOB: &str = "did:plc:bbbbbbbbbbbbbbbbbbbbbbbb";

    fn did(s: &str) -> Did {
        Did::new_owned(s).unwrap()
    }

    /// Build a minimal valid session. The DPoP key is minted via the public
    /// keyset API and moved straight into the holder without naming `jose_jwk`;
    /// its contents are never asserted (the store never validates it).
    fn client_session(
        account_did: &str,
        session_id: &str,
        access_token: &str,
    ) -> ClientSessionData {
        let dpop_key = Keyset::generate_es256("test")
            .unwrap()
            .public_jwks()
            .keys
            .into_iter()
            .next()
            .unwrap()
            .key;
        ClientSessionData {
            account_did: did(account_did),
            session_id: SmolStr::new(session_id),
            host_url: Uri::parse("https://pds.example.com").unwrap().to_owned(),
            authserver_url: SmolStr::new("https://issuer.example.com"),
            authserver_token_endpoint: SmolStr::new("https://issuer.example.com/token"),
            authserver_revocation_endpoint: None,
            scopes: Scopes::empty(),
            dpop_data: DpopClientData {
                dpop_key,
                dpop_authserver_nonce: SmolStr::default(),
                dpop_host_nonce: SmolStr::default(),
            },
            token_set: TokenSet {
                iss: SmolStr::new("https://issuer.example.com"),
                sub: did(account_did),
                aud: SmolStr::new("https://pds.example.com"),
                scope: None,
                refresh_token: None,
                access_token: SmolStr::new(access_token),
                token_type: OAuthTokenType::DPoP,
                expires_at: None,
            },
            resolved_scopes: None,
        }
    }

    fn auth_request(state: &str, account_did: Option<&str>) -> AuthRequestData {
        let dpop_key = Keyset::generate_es256("test")
            .unwrap()
            .public_jwks()
            .keys
            .into_iter()
            .next()
            .unwrap()
            .key;
        AuthRequestData {
            state: SmolStr::new(state),
            authserver_url: SmolStr::new("https://issuer.example.com"),
            account_did: account_did.map(did),
            scopes: Scopes::empty(),
            request_uri: SmolStr::new("urn:ietf:params:oauth:request_uri:abc"),
            authserver_token_endpoint: SmolStr::new("https://issuer.example.com/token"),
            authserver_revocation_endpoint: None,
            pkce_verifier: SmolStr::new("verifier"),
            dpop_data: DpopReqData {
                dpop_key,
                dpop_authserver_nonce: None,
            },
        }
    }

    async fn store() -> SqliteAuthStore {
        let options = SqliteConnectOptions::from_str("sqlite::memory:").unwrap();
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect_with(options)
            .await
            .unwrap();
        sqlx::migrate!("./migrations").run(&pool).await.unwrap();
        SqliteAuthStore::new(pool)
    }

    #[tokio::test]
    async fn upsert_get_roundtrip_returns_the_stored_session() {
        let store = store().await;
        assert!(
            store
                .get_session(&did(ALICE), "s1")
                .await
                .unwrap()
                .is_none()
        );

        store
            .upsert_session(client_session(ALICE, "s1", "tok"))
            .await
            .unwrap();

        let got = store.get_session(&did(ALICE), "s1").await.unwrap().unwrap();
        assert_eq!(got.account_did, did(ALICE));
        assert_eq!(got.session_id.as_str(), "s1");
        assert_eq!(got.token_set.access_token.as_str(), "tok");
    }

    #[tokio::test]
    async fn upsert_overwrites_same_key_latest_wins() {
        let store = store().await;
        store
            .upsert_session(client_session(ALICE, "s1", "first"))
            .await
            .unwrap();
        let first_created_at: i64 = sqlx::query_scalar(
            "SELECT created_at FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
        )
        .bind(ALICE)
        .bind("s1")
        .fetch_one(&store.pool)
        .await
        .unwrap();
        let first_updated_at: i64 = sqlx::query_scalar(
            "SELECT updated_at FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
        )
        .bind(ALICE)
        .bind("s1")
        .fetch_one(&store.pool)
        .await
        .unwrap();

        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        store
            .upsert_session(client_session(ALICE, "s1", "second"))
            .await
            .unwrap();
        let second_created_at: i64 = sqlx::query_scalar(
            "SELECT created_at FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
        )
        .bind(ALICE)
        .bind("s1")
        .fetch_one(&store.pool)
        .await
        .unwrap();
        let second_updated_at: i64 = sqlx::query_scalar(
            "SELECT updated_at FROM oauth_sessions WHERE account_did = ? AND session_id = ?",
        )
        .bind(ALICE)
        .bind("s1")
        .fetch_one(&store.pool)
        .await
        .unwrap();

        assert_eq!(
            second_created_at, first_created_at,
            "upsert must preserve original created_at"
        );
        assert!(
            second_updated_at > first_updated_at,
            "upsert must advance updated_at"
        );

        // Exactly one row for the key (no duplicate), and the payload reflects
        // the latest write — proving upsert-overwrite + payload round-trip.
        assert_eq!(store.list_session_keys().await.unwrap().len(), 1);
        let got = store.get_session(&did(ALICE), "s1").await.unwrap().unwrap();
        assert_eq!(got.token_set.access_token.as_str(), "second");
    }

    #[tokio::test]
    async fn delete_removes_only_the_target_session() {
        let store = store().await;
        store
            .upsert_session(client_session(ALICE, "s1", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(ALICE, "s2", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(BOB, "s3", "t"))
            .await
            .unwrap();

        store.delete_session(&did(ALICE), "s1").await.unwrap();

        let keys = store.list_session_keys().await.unwrap();
        assert_eq!(keys.len(), 2);
        assert!(
            store
                .get_session(&did(ALICE), "s1")
                .await
                .unwrap()
                .is_none()
        );
        assert!(
            store
                .get_session(&did(ALICE), "s2")
                .await
                .unwrap()
                .is_some()
        );
        assert!(store.get_session(&did(BOB), "s3").await.unwrap().is_some());
    }

    #[tokio::test]
    async fn list_session_keys_enumerates_and_is_filterable_by_did() {
        let store = store().await;
        store
            .upsert_session(client_session(ALICE, "s1", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(ALICE, "s2", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(BOB, "s3", "t"))
            .await
            .unwrap();

        // The structured `account_did` column is queryable *without* loading the
        // payload: a direct COUNT groups Alice's sessions by DID alone.
        let alice_count: i64 =
            sqlx::query_scalar("SELECT COUNT(*) FROM oauth_sessions WHERE account_did = ?")
                .bind(ALICE)
                .fetch_one(&store.pool)
                .await
                .unwrap();
        assert_eq!(alice_count, 2);

        let keys = store.list_session_keys().await.unwrap();
        assert_eq!(keys.len(), 3);
        assert!(
            keys.iter()
                .any(|k| k.did == did(ALICE) && k.session_id == "s1")
        );
        assert!(
            keys.iter()
                .any(|k| k.did == did(ALICE) && k.session_id == "s2")
        );
        assert!(
            keys.iter()
                .any(|k| k.did == did(BOB) && k.session_id == "s3")
        );
    }

    #[tokio::test]
    async fn select_session_resolves_key_did_and_any() {
        let store = store().await;
        store
            .upsert_session(client_session(ALICE, "s1", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(ALICE, "s2", "t"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(BOB, "s3", "t"))
            .await
            .unwrap();

        // Exact key.
        let m = store
            .select_session(&SessionHint::key(SessionKey::new(did(ALICE), "s2")))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(m.key.session_id.as_str(), "s2");

        // By DID — must be one of Alice's, never Bob's.
        let m = store
            .select_session(&SessionHint::did(did(ALICE)))
            .await
            .unwrap()
            .unwrap();
        assert_eq!(m.key.did, did(ALICE));

        // Any — some session.
        let m = store
            .select_session(&SessionHint::any())
            .await
            .unwrap()
            .unwrap();
        assert!(
            store
                .get_session(&m.key.did, m.key.session_id.as_str())
                .await
                .unwrap()
                .is_some()
        );

        // Handle/identifier hints never resolve against a resolver-free store.
        assert!(
            store
                .select_session(&SessionHint::identifier("alice.example.com".into()))
                .await
                .unwrap()
                .is_none()
        );
    }

    #[tokio::test]
    async fn select_session_returns_none_when_no_match() {
        let store = store().await;
        assert!(
            store
                .select_session(&SessionHint::any())
                .await
                .unwrap()
                .is_none()
        );
        assert!(
            store
                .select_session(&SessionHint::key(SessionKey::new(did(ALICE), "x")))
                .await
                .unwrap()
                .is_none()
        );
        assert!(
            store
                .select_session(&SessionHint::did(did(ALICE)))
                .await
                .unwrap()
                .is_none()
        );

        store
            .upsert_session(client_session(BOB, "s3", "t"))
            .await
            .unwrap();
        // A DID with no session still resolves to None even when others exist.
        assert!(
            store
                .select_session(&SessionHint::did(did(ALICE)))
                .await
                .unwrap()
                .is_none()
        );
    }

    #[tokio::test]
    async fn auth_request_lifecycle_save_get_delete() {
        let store = store().await;
        assert!(store.get_auth_req_info("state-1").await.unwrap().is_none());

        store
            .save_auth_req_info(&auth_request("state-1", Some(ALICE)))
            .await
            .unwrap();
        let got = store.get_auth_req_info("state-1").await.unwrap().unwrap();
        assert_eq!(got.state.as_str(), "state-1");
        assert_eq!(got.account_did.as_ref().unwrap(), &did(ALICE));

        store.delete_auth_req_info("state-1").await.unwrap();
        assert!(store.get_auth_req_info("state-1").await.unwrap().is_none());

        // A second request is unaffected by deleting the first.
        store
            .save_auth_req_info(&auth_request("state-2", None))
            .await
            .unwrap();
        store.delete_auth_req_info("state-1").await.unwrap();
        assert!(store.get_auth_req_info("state-2").await.unwrap().is_some());
    }

    #[tokio::test]
    async fn cross_session_isolation() {
        let store = store().await;
        store
            .upsert_session(client_session(ALICE, "s1", "a1"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(ALICE, "s2", "a2"))
            .await
            .unwrap();
        store
            .upsert_session(client_session(BOB, "s1", "b1"))
            .await
            .unwrap();

        // Two sessions share session_id "s1" under different DIDs.
        assert_eq!(
            store
                .get_session(&did(ALICE), "s1")
                .await
                .unwrap()
                .unwrap()
                .token_set
                .access_token
                .as_str(),
            "a1"
        );
        assert_eq!(
            store
                .get_session(&did(BOB), "s1")
                .await
                .unwrap()
                .unwrap()
                .token_set
                .access_token
                .as_str(),
            "b1"
        );
    }
}