jacquard 0.12.0-beta.2

use jacquard_common::deps::fluent_uri::Uri;
use jacquard_common::session::{FileTokenStore, SessionStore, SessionStoreError};
use jacquard_common::types::string::{Datetime, Did};
use jacquard_oauth::scopes::Scopes;
use jacquard_oauth::session::{AuthRequestData, ClientSessionData, DpopClientData, DpopReqData};
use jacquard_oauth::types::OAuthTokenType;
use jose_jwk::Key;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use smol_str::SmolStr;

/// On-disk session records for app-password and OAuth flows, sharing a single JSON map.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum StoredSession {
    /// App-password session
    Atp(StoredAtSession),
    /// OAuth client session
    OAuth(OAuthSession),
    /// OAuth authorization request state
    OAuthState(OAuthState),
}

/// Minimal persisted representation of an app‑password session.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredAtSession {
    /// Access token (JWT)
    access_jwt: String,
    /// Refresh token (JWT)
    refresh_jwt: String,
    /// Account DID
    did: String,
    /// Optional PDS endpoint for faster resume
    #[serde(skip_serializing_if = "std::option::Option::is_none")]
    pds: Option<String>,
    /// Session id label (e.g., "session")
    session_id: String,
    /// Last known handle
    handle: String,
}

/// Persisted OAuth client session (on-disk format).
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct OAuthSession {
    /// Account DID
    account_did: String,
    /// Client-generated session id (usually auth `state`)
    session_id: String,

    /// Base URL of the resource server (PDS)
    host_url: Uri<String>,

    /// Base URL of the authorization server (PDS or entryway)
    authserver_url: String,

    /// Full token endpoint URL
    authserver_token_endpoint: String,

    /// Full revocation endpoint URL, if available
    #[serde(skip_serializing_if = "std::option::Option::is_none")]
    authserver_revocation_endpoint: Option<String>,

    /// Granted scopes (space-separated, normalized).
    scopes: String,

    /// Client DPoP key material
    pub dpop_key: Key,
    /// Current auth server DPoP nonce
    pub dpop_authserver_nonce: String,
    /// Current resource server (PDS) DPoP nonce
    pub dpop_host_nonce: String,

    /// Token response issuer
    pub iss: String,
    /// Token subject (DID)
    pub sub: String,
    /// Token audience (verified PDS URL)
    pub aud: String,
    /// Token scopes (raw) if provided
    pub scope: Option<String>,

    /// Refresh token
    pub refresh_token: Option<String>,
    /// Access token
    pub access_token: String,
    /// Token type (e.g., DPoP)
    pub token_type: OAuthTokenType,

    /// Expiration timestamp
    pub expires_at: Option<Datetime>,
}

impl<S: jacquard_common::bos::BosStr + Ord> From<ClientSessionData<S>> for OAuthSession {
    fn from(data: ClientSessionData<S>) -> Self {
        OAuthSession {
            account_did: AsRef::<str>::as_ref(&data.account_did).to_owned(),
            session_id: AsRef::<str>::as_ref(&data.session_id).to_owned(),
            host_url: data.host_url.clone(),
            authserver_url: AsRef::<str>::as_ref(&data.authserver_url).to_owned(),
            authserver_token_endpoint: AsRef::<str>::as_ref(&data.authserver_token_endpoint)
                .to_owned(),
            authserver_revocation_endpoint: data
                .authserver_revocation_endpoint
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            scopes: String::from(data.scopes.to_normalized_string()),
            dpop_key: data.dpop_data.dpop_key,
            dpop_authserver_nonce: AsRef::<str>::as_ref(&data.dpop_data.dpop_authserver_nonce)
                .to_owned(),
            dpop_host_nonce: AsRef::<str>::as_ref(&data.dpop_data.dpop_host_nonce).to_owned(),
            iss: AsRef::<str>::as_ref(&data.token_set.iss).to_owned(),
            sub: AsRef::<str>::as_ref(&data.token_set.sub).to_owned(),
            aud: AsRef::<str>::as_ref(&data.token_set.aud).to_owned(),
            scope: data
                .token_set
                .scope
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            refresh_token: data
                .token_set
                .refresh_token
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            access_token: AsRef::<str>::as_ref(&data.token_set.access_token).to_owned(),
            token_type: data.token_set.token_type,
            expires_at: data.token_set.expires_at,
        }
    }
}

impl From<OAuthSession> for ClientSessionData {
    fn from(session: OAuthSession) -> Self {
        ClientSessionData {
            account_did: Did::new_owned(session.account_did).expect("stored DID should be valid"),
            session_id: SmolStr::from(session.session_id),
            host_url: session.host_url,
            authserver_url: SmolStr::from(session.authserver_url),
            authserver_token_endpoint: SmolStr::from(session.authserver_token_endpoint),
            authserver_revocation_endpoint: session
                .authserver_revocation_endpoint
                .map(SmolStr::from),
            scopes: Scopes::new(SmolStr::from(session.scopes.as_str()))
                .expect("stored scopes should be valid"),
            dpop_data: DpopClientData {
                dpop_key: session.dpop_key,
                dpop_authserver_nonce: SmolStr::from(session.dpop_authserver_nonce),
                dpop_host_nonce: SmolStr::from(session.dpop_host_nonce),
            },
            token_set: jacquard_oauth::types::TokenSet {
                iss: SmolStr::from(session.iss),
                sub: Did::new_owned(session.sub).expect("stored DID should be valid"),
                aud: SmolStr::from(session.aud),
                scope: session.scope.map(SmolStr::from),
                refresh_token: session.refresh_token.map(SmolStr::from),
                access_token: SmolStr::from(session.access_token),
                token_type: session.token_type,
                expires_at: session.expires_at,
            },
            #[cfg(feature = "scope-check")]
            resolved_scopes: None,
        }
    }
}

/// Persisted OAuth authorization request state.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct OAuthState {
    /// Random identifier generated for the authorization flow (`state`)
    pub state: String,

    /// Base URL of the authorization server (PDS or entryway)
    pub authserver_url: Uri<String>,

    /// Optional pre-known account DID
    #[serde(skip_serializing_if = "std::option::Option::is_none")]
    pub account_did: Option<String>,

    /// Requested scopes (space-separated, normalized).
    pub scopes: String,

    /// Request URI for the authorization step
    pub request_uri: String,

    /// Full token endpoint URL
    pub authserver_token_endpoint: String,

    /// Full revocation endpoint URL, if available
    #[serde(skip_serializing_if = "std::option::Option::is_none")]
    pub authserver_revocation_endpoint: Option<String>,

    /// PKCE verifier
    pub pkce_verifier: String,

    /// Client DPoP key material
    pub dpop_key: Key,
    /// Auth server DPoP nonce at PAR time
    #[serde(skip_serializing_if = "std::option::Option::is_none")]
    pub dpop_authserver_nonce: Option<String>,
}

impl<S: jacquard_common::bos::BosStr + Ord> TryFrom<AuthRequestData<S>> for OAuthState {
    type Error = jacquard_common::deps::fluent_uri::ParseError;

    fn try_from(value: AuthRequestData<S>) -> Result<Self, Self::Error> {
        Ok(OAuthState {
            authserver_url: Uri::parse(value.authserver_url.as_ref())?.to_owned(),
            account_did: value
                .account_did
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            scopes: String::from(value.scopes.to_normalized_string()),
            request_uri: AsRef::<str>::as_ref(&value.request_uri).to_owned(),
            authserver_token_endpoint: AsRef::<str>::as_ref(&value.authserver_token_endpoint)
                .to_owned(),
            authserver_revocation_endpoint: value
                .authserver_revocation_endpoint
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            pkce_verifier: AsRef::<str>::as_ref(&value.pkce_verifier).to_owned(),
            dpop_key: value.dpop_data.dpop_key,
            dpop_authserver_nonce: value
                .dpop_data
                .dpop_authserver_nonce
                .map(|s| AsRef::<str>::as_ref(&s).to_owned()),
            state: AsRef::<str>::as_ref(&value.state).to_owned(),
        })
    }
}

impl From<OAuthState> for AuthRequestData {
    fn from(value: OAuthState) -> Self {
        AuthRequestData {
            authserver_url: SmolStr::from(value.authserver_url.as_str()),
            state: SmolStr::from(value.state),
            account_did: value
                .account_did
                .map(|s| Did::new_owned(s).expect("stored DID should be valid")),
            authserver_revocation_endpoint: value.authserver_revocation_endpoint.map(SmolStr::from),
            scopes: Scopes::new(SmolStr::from(value.scopes.as_str()))
                .expect("stored scopes should be valid"),
            request_uri: SmolStr::from(value.request_uri),
            authserver_token_endpoint: SmolStr::from(value.authserver_token_endpoint),
            pkce_verifier: SmolStr::from(value.pkce_verifier),
            dpop_data: DpopReqData {
                dpop_key: value.dpop_key,
                dpop_authserver_nonce: value.dpop_authserver_nonce.map(SmolStr::from),
            },
        }
    }
}

/// Convenience wrapper over `FileTokenStore` offering unified storage across auth modes.
pub struct FileAuthStore(FileTokenStore);

impl FileAuthStore {
    /// Create a new file-backed auth store wrapping `FileTokenStore`.
    ///
    /// # Errors
    ///
    /// Returns an error if parent directories cannot be created or the file cannot be written.
    pub fn try_new(path: impl AsRef<std::path::Path>) -> Result<Self, SessionStoreError> {
        Ok(Self(FileTokenStore::try_new(path)?))
    }

    /// Create a new file-backed auth store wrapping `FileTokenStore`.
    ///
    /// # Panics
    ///
    /// Panics if parent directories cannot be created or the file cannot be written.
    /// Prefer [`try_new`](Self::try_new) for fallible construction.
    pub fn new(path: impl AsRef<std::path::Path>) -> Self {
        Self(FileTokenStore::new(path))
    }
}

impl jacquard_oauth::authstore::ClientAuthStore for FileAuthStore {
    async fn get_session<D: jacquard_common::bos::BosStr + Send + Sync>(
        &self,
        did: &Did<D>,
        session_id: &str,
    ) -> Result<Option<ClientSessionData>, SessionStoreError> {
        let key = format!("{}_{}", did, session_id);
        if let StoredSession::OAuth(session) = self
            .0
            .get(&key)
            .await
            .ok_or(SessionStoreError::Other("not found".into()))?
        {
            Ok(Some(session.into()))
        } else {
            Ok(None)
        }
    }

    async fn upsert_session(&self, session: ClientSessionData) -> Result<(), SessionStoreError> {
        let key = format!("{}_{}", session.account_did, session.session_id);
        self.0
            .set(key, StoredSession::OAuth(session.into()))
            .await?;
        Ok(())
    }

    async fn delete_session<D: jacquard_common::bos::BosStr + Send + Sync>(
        &self,
        did: &Did<D>,
        session_id: &str,
    ) -> Result<(), SessionStoreError> {
        let key = format!("{}_{}", did, session_id);
        let file = std::fs::read_to_string(&self.0.path)?;
        let mut store: Value = serde_json::from_str(&file)?;
        let key_string = key.to_string();
        if let Some(store) = store.as_object_mut() {
            store.remove(&key_string);

            std::fs::write(&self.0.path, serde_json::to_string_pretty(&store)?)?;
            Ok(())
        } else {
            Err(SessionStoreError::Other("invalid store".into()))
        }
    }

    async fn get_auth_req_info(
        &self,
        state: &str,
    ) -> Result<Option<AuthRequestData>, SessionStoreError> {
        let key = format!("authreq_{}", state);
        if let StoredSession::OAuthState(auth_req) = self
            .0
            .get(&key)
            .await
            .ok_or(SessionStoreError::Other("not found".into()))?
        {
            Ok(Some(auth_req.into()))
        } else {
            Ok(None)
        }
    }

    async fn save_auth_req_info(
        &self,
        auth_req_info: &AuthRequestData,
    ) -> Result<(), SessionStoreError> {
        let key = format!("authreq_{}", auth_req_info.state);
        let state = auth_req_info.clone().try_into().map_err(
            |e: jacquard_common::deps::fluent_uri::ParseError| {
                SessionStoreError::Other(Box::new(e))
            },
        )?;
        self.0.set(key, StoredSession::OAuthState(state)).await?;
        Ok(())
    }

    async fn delete_auth_req_info(&self, state: &str) -> Result<(), SessionStoreError> {
        let key = format!("authreq_{}", state);
        let file = std::fs::read_to_string(&self.0.path)?;
        let mut store: Value = serde_json::from_str(&file)?;
        let key_string = key.to_string();
        if let Some(store) = store.as_object_mut() {
            store.remove(&key_string);

            std::fs::write(&self.0.path, serde_json::to_string_pretty(&store)?)?;
            Ok(())
        } else {
            Err(SessionStoreError::Other("invalid store".into()))
        }
    }
}

impl FileAuthStore {
    /// Update the persisted PDS endpoint for an app-password session (best-effort).
    pub fn set_atp_pds(
        &self,
        key: &crate::client::credential_session::SessionKey,
        pds: &Uri<String>,
    ) -> Result<(), SessionStoreError> {
        let key_str = format!("{}_{}", key.0, key.1);
        let file = std::fs::read_to_string(&self.0.path)?;
        let mut store: Value = serde_json::from_str(&file)?;
        if let Some(map) = store.as_object_mut() {
            if let Some(value) = map.get_mut(&key_str) {
                if let Some(outer) = value.as_object_mut() {
                    if let Some(inner) = outer.get_mut("Atp").and_then(|v| v.as_object_mut()) {
                        inner.insert(
                            "pds".to_string(),
                            serde_json::Value::String(pds.as_str().to_string()),
                        );
                        std::fs::write(&self.0.path, serde_json::to_string_pretty(&store)?)?;
                        return Ok(());
                    }
                }
            }
        }
        Err(SessionStoreError::Other("invalid store".into()))
    }

    /// Read the persisted PDS endpoint for an app-password session, if present.
    pub fn get_atp_pds(
        &self,
        key: &crate::client::credential_session::SessionKey,
    ) -> Result<Option<Uri<String>>, SessionStoreError> {
        let key_str = format!("{}_{}", key.0, key.1);
        let file = std::fs::read_to_string(&self.0.path)?;
        let store: Value = serde_json::from_str(&file)?;
        if let Some(value) = store.get(&key_str) {
            if let Some(obj) = value.as_object() {
                if let Some(serde_json::Value::Object(inner)) = obj.get("Atp") {
                    if let Some(serde_json::Value::String(pds)) = inner.get("pds") {
                        return Ok(Uri::parse(pds.as_str()).ok().map(|u| u.to_owned()));
                    }
                }
            }
        }
        Ok(None)
    }
}

impl
    jacquard_common::session::SessionStore<
        crate::client::credential_session::SessionKey,
        crate::client::AtpSession,
    > for FileAuthStore
{
    async fn get(
        &self,
        key: &crate::client::credential_session::SessionKey,
    ) -> Option<crate::client::AtpSession> {
        let key_str = format!("{}_{}", key.0, key.1);
        if let Some(StoredSession::Atp(stored)) = self.0.get(&key_str).await {
            Some(crate::client::AtpSession {
                access_jwt: stored.access_jwt.into(),
                refresh_jwt: stored.refresh_jwt.into(),
                did: stored.did.into(),
                handle: stored.handle.into(),
            })
        } else {
            None
        }
    }

    async fn set(
        &self,
        key: crate::client::credential_session::SessionKey,
        session: crate::client::AtpSession,
    ) -> Result<(), jacquard_common::session::SessionStoreError> {
        let key_str = format!("{}_{}", key.0, key.1);
        let stored = StoredAtSession {
            access_jwt: session.access_jwt.to_string(),
            refresh_jwt: session.refresh_jwt.to_string(),
            did: session.did.to_string(),
            // pds endpoint is resolved on restore; do not persist
            pds: None,
            session_id: key.1.to_string(),
            handle: session.handle.to_string(),
        };
        self.0.set(key_str, StoredSession::Atp(stored)).await
    }

    async fn del(
        &self,
        key: &crate::client::credential_session::SessionKey,
    ) -> Result<(), jacquard_common::session::SessionStoreError> {
        let key_str = format!("{}_{}", key.0, key.1);
        // Manual removal to mirror existing pattern
        let file = std::fs::read_to_string(&self.0.path)?;
        let mut store: serde_json::Value = serde_json::from_str(&file)?;
        if let Some(map) = store.as_object_mut() {
            map.remove(&key_str);
            std::fs::write(&self.0.path, serde_json::to_string_pretty(&store)?)?;
            Ok(())
        } else {
            Err(jacquard_common::session::SessionStoreError::Other(
                "invalid store".into(),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client::AtpSession;
    use crate::client::credential_session::SessionKey;
    use jacquard_common::types::string::{Did, Handle};
    use std::fs;
    use std::path::PathBuf;

    fn temp_file() -> PathBuf {
        let mut p = std::env::temp_dir();
        p.push(format!("jacquard-test-{}.json", std::process::id()));
        p
    }

    #[tokio::test]
    async fn file_auth_store_roundtrip_atp() {
        let path = temp_file();
        // initialize empty store file
        fs::write(&path, "{}").unwrap();
        let store = FileAuthStore::new(&path);
        let session = AtpSession {
            access_jwt: "a".into(),
            refresh_jwt: "r".into(),
            did: Did::new_static("did:plc:alice").unwrap(),
            handle: Handle::new_static("alice.bsky.social").unwrap(),
        };
        let key = SessionKey(session.did.clone(), "session".into());
        jacquard_common::session::SessionStore::set(&store, key.clone(), session.clone())
            .await
            .unwrap();
        let restored = jacquard_common::session::SessionStore::get(&store, &key)
            .await
            .unwrap();
        assert_eq!(restored.access_jwt.as_str(), "a");
        // clean up
        let _ = fs::remove_file(&path);
    }
}