vta-service 0.2.0

use axum::Json;
use axum::extract::{Path, Query, State};
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use vta_sdk::protocols::auth::{
    AuthenticateData, AuthenticateResponse, ChallengeData, ChallengeRequest, ChallengeResponse,
};
use vta_sdk::protocols::credential_management::generate::GenerateCredentialsResultBody;

use crate::acl::{Role, check_acl, check_acl_full};
use crate::audit::audit;
use crate::auth::{AdminAuth, AuthClaims, ManageAuth};
use crate::auth::session::{
    Session, SessionState, delete_session, get_session, get_session_by_refresh, list_sessions,
    now_epoch, store_refresh_index, store_session, update_session,
};
#[cfg(feature = "tee")]
use crate::error::tee_attestation_error;
use crate::error::AppError;
use crate::operations;
use crate::server::AppState;
#[cfg(feature = "tee")]
use tracing::error;
use tracing::{info, warn};

// ---------- POST /auth/challenge ----------

/// POST /auth/challenge — issue a DID-auth challenge nonce for a session. Auth: unauthenticated.
pub async fn challenge(
    State(state): State<AppState>,
    Json(req): Json<ChallengeRequest>,
) -> Result<Json<ChallengeResponse>, AppError> {
    // DID method whitelist enforcement (TEE mode)
    #[cfg(feature = "tee")]
    {
        let config = state.config.read().await;
        if let Some(ref allowed) = config.tee.allowed_did_methods {
            let did_ok = allowed.iter().any(|prefix| req.did.starts_with(prefix));
            if !did_ok {
                warn!(did = %req.did, "auth rejected: DID method not in allowed_did_methods");
                return Err(AppError::Forbidden(format!(
                    "DID method not allowed — accepted methods: {}",
                    allowed.join(", ")
                )));
            }
        }
        drop(config);
    }

    // ACL enforcement: DID must be in the ACL to request a challenge
    check_acl(&state.acl_ks, &req.did).await?;

    let session_id = Uuid::new_v4().to_string();

    // Generate 32-byte random challenge as hex
    let mut challenge_bytes = [0u8; 32];
    rand::fill(&mut challenge_bytes);
    let mut challenge = hex::encode(challenge_bytes);

    // Nonce replay prevention: store the challenge hash to detect reuse.
    // Challenges are random 32 bytes so collision is negligible, but this
    // provides defense in depth against replay attacks.
    let nonce_key = format!("nonce:{challenge}");
    if state.sessions_ks.get_raw(nonce_key.clone()).await?.is_some() {
        warn!(challenge = %challenge, "challenge nonce collision detected — regenerating");
        // Extremely unlikely (2^-256) but handle gracefully
        rand::fill(&mut challenge_bytes);
        challenge = hex::encode(challenge_bytes);
    }
    state.sessions_ks
        .insert_raw(format!("nonce:{challenge}"), session_id.as_bytes().to_vec())
        .await?;

    let session = Session {
        session_id: session_id.clone(),
        did: req.did,
        challenge: challenge.clone(),
        state: SessionState::ChallengeSent,
        created_at: now_epoch(),
        refresh_token: None,
        refresh_expires_at: None,
    };

    store_session(&state.sessions_ks, &session).await?;

    // Optionally bind a TEE attestation report to the challenge nonce.
    // This proves the challenge was generated inside a trusted execution environment.
    #[cfg(feature = "tee")]
    let tee_attestation = if let Some(ref tee) = state.tee {
        let config = state.config.read().await;
        let vta_did = config.vta_did.clone();
        drop(config);

        let user_data = vta_did.as_deref().unwrap_or("").as_bytes();
        let nonce_bytes = &challenge_bytes[..];

        match tee.state.provider.attest(user_data, nonce_bytes) {
            Ok(mut report) => {
                report.vta_did = vta_did;
                Some(serde_json::to_value(&report).map_err(|e| {
                    AppError::Internal(format!("failed to serialize attestation report: {e}"))
                })?)
            }
            Err(e) => {
                // In TEE required mode, attestation failure is a hard error.
                // A broken TEE must not silently serve unattested challenges.
                let tee_mode = &state.config.read().await.tee.mode;
                if *tee_mode == crate::config::TeeMode::Required {
                    error!("TEE attestation failed in required mode — refusing challenge: {e}");
                    return Err(tee_attestation_error(format!(
                        "TEE attestation failed (mode=required): {e}"
                    )));
                }
                warn!("TEE attestation failed (mode=optional) — challenge served without attestation: {e}");
                None
            }
        }
    } else {
        None
    };
    #[cfg(not(feature = "tee"))]
    let tee_attestation = None;

    info!(did = %session.did, session_id = %session.session_id, "auth challenge issued");
    audit!("auth.challenge", actor = &session.did, resource =&session.session_id, outcome = "success");

    Ok(Json(ChallengeResponse {
        session_id,
        data: ChallengeData {
            challenge,
            tee_attestation,
        },
    }))
}

// ---------- POST /auth/ ----------

/// POST /auth/ — verify a signed DIDComm challenge and issue access+refresh tokens. Auth: unauthenticated.
pub async fn authenticate(
    State(state): State<AppState>,
    body: String,
) -> Result<Json<AuthenticateResponse>, AppError> {
    let atm = state
        .atm
        .as_ref()
        .ok_or_else(|| AppError::Authentication("ATM not configured".into()))?;
    let jwt_keys = state
        .jwt_keys
        .as_ref()
        .ok_or_else(|| AppError::Authentication("JWT keys not configured".into()))?;

    // Unpack the DIDComm message
    let (msg, _metadata) = atm
        .unpack(&body)
        .await
        .map_err(|e| AppError::Authentication(format!("failed to unpack message: {e}")))?;

    // Validate message type
    if msg.typ != "https://affinidi.com/atm/1.0/authenticate" {
        return Err(AppError::Authentication(format!(
            "unexpected message type: {}",
            msg.typ
        )));
    }

    // Extract challenge and session_id from body
    let challenge = msg.body["challenge"]
        .as_str()
        .ok_or_else(|| AppError::Authentication("missing challenge in message body".into()))?;
    let session_id = msg.body["session_id"]
        .as_str()
        .ok_or_else(|| AppError::Authentication("missing session_id in message body".into()))?;

    // Validate sender DID
    let sender_did = msg
        .from
        .as_deref()
        .ok_or_else(|| AppError::Authentication("message has no sender (from)".into()))?;

    // Look up session and validate
    let mut session = get_session(&state.sessions_ks, session_id)
        .await?
        .ok_or_else(|| AppError::Authentication("session not found".into()))?;

    if session.state != SessionState::ChallengeSent {
        warn!(session_id, "authentication rejected: session replay");
        audit!("auth.authenticate", actor = sender_did.split('#').next().unwrap_or(sender_did), resource =session_id, outcome = "denied:replay");
        return Err(AppError::Authentication(
            "session already authenticated (replay)".into(),
        ));
    }
    if session.challenge != challenge {
        warn!(session_id, "authentication rejected: challenge mismatch");
        audit!("auth.authenticate", actor = sender_did.split('#').next().unwrap_or(sender_did), resource =session_id, outcome = "denied:challenge_mismatch");
        return Err(AppError::Authentication("challenge mismatch".into()));
    }
    // Match the DID (compare base DID, ignoring any fragment)
    let sender_base = sender_did.split('#').next().unwrap_or(sender_did);
    if session.did != sender_base {
        warn!(session_id, sender = %sender_base, expected = %session.did, "authentication rejected: DID mismatch");
        audit!("auth.authenticate", actor = sender_base, resource =session_id, outcome = "denied:did_mismatch");
        return Err(AppError::Authentication("DID mismatch".into()));
    }

    // Read all auth config values in a single lock acquisition
    let (challenge_ttl, access_expiry, refresh_expiry) = {
        let config = state.config.read().await;
        (config.auth.challenge_ttl, config.auth.access_token_expiry, config.auth.refresh_token_expiry)
    };

    // Check challenge TTL
    if now_epoch().saturating_sub(session.created_at) > challenge_ttl {
        warn!(session_id, "authentication rejected: challenge expired");
        audit!("auth.authenticate", actor = sender_base, resource =session_id, outcome = "denied:expired");
        return Err(AppError::Authentication("challenge expired".into()));
    }

    // Look up ACL entry to get role and allowed contexts for the token
    let (role, allowed_contexts) = check_acl_full(&state.acl_ks, &session.did).await?;

    // Check if VTA is running in a TEE
    #[cfg(feature = "tee")]
    let tee_attested = state.tee.is_some();
    #[cfg(not(feature = "tee"))]
    let tee_attested = false;

    let claims = jwt_keys.new_claims(
        session.did.clone(),
        session.session_id.clone(),
        role.to_string(),
        allowed_contexts,
        access_expiry,
        tee_attested,
    );
    let access_expires_at = claims.exp;
    let access_token = jwt_keys.encode(&claims)?;

    let refresh_token = Uuid::new_v4().to_string();
    let refresh_expires_at = now_epoch() + refresh_expiry;

    // Update session to Authenticated
    session.state = SessionState::Authenticated;
    session.refresh_token = Some(refresh_token.clone());
    session.refresh_expires_at = Some(refresh_expires_at);
    update_session(&state.sessions_ks, &session).await?;

    // Store reverse refresh index
    store_refresh_index(&state.sessions_ks, &refresh_token, &session.session_id).await?;

    info!(did = %session.did, session_id = %session.session_id, "authentication successful");
    audit!("auth.authenticate", actor = &session.did, resource =&session.session_id, outcome = "success");

    Ok(Json(AuthenticateResponse {
        session_id: Some(session.session_id),
        data: AuthenticateData {
            access_token,
            access_expires_at,
            refresh_token: Some(refresh_token),
            refresh_expires_at: Some(refresh_expires_at),
        },
    }))
}

// ---------- POST /auth/refresh ----------

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RefreshResponse {
    pub session_id: String,
    pub data: RefreshData,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RefreshData {
    pub access_token: String,
    pub access_expires_at: u64,
}

/// POST /auth/refresh — exchange a refresh token for a new access token. Auth: unauthenticated.
pub async fn refresh(
    State(state): State<AppState>,
    body: String,
) -> Result<Json<RefreshResponse>, AppError> {
    let atm = state
        .atm
        .as_ref()
        .ok_or_else(|| AppError::Authentication("ATM not configured".into()))?;
    let jwt_keys = state
        .jwt_keys
        .as_ref()
        .ok_or_else(|| AppError::Authentication("JWT keys not configured".into()))?;

    // Unpack the DIDComm message
    let (msg, _metadata) = atm
        .unpack(&body)
        .await
        .map_err(|e| AppError::Authentication(format!("failed to unpack message: {e}")))?;

    // Validate message type
    if msg.typ != "https://affinidi.com/atm/1.0/authenticate/refresh" {
        return Err(AppError::Authentication(format!(
            "unexpected message type: {}",
            msg.typ
        )));
    }

    // Extract refresh_token from body
    let refresh_token = msg.body["refresh_token"]
        .as_str()
        .ok_or_else(|| AppError::Authentication("missing refresh_token in message body".into()))?;

    // Look up session by refresh token
    let session_id = get_session_by_refresh(&state.sessions_ks, refresh_token)
        .await?
        .ok_or_else(|| AppError::Authentication("refresh token not found".into()))?;

    let session = get_session(&state.sessions_ks, &session_id)
        .await?
        .ok_or_else(|| AppError::Authentication("session not found".into()))?;

    if session.state != SessionState::Authenticated {
        return Err(AppError::Authentication("session not authenticated".into()));
    }

    // Verify refresh token hasn't expired
    if let Some(expires_at) = session.refresh_expires_at
        && now_epoch() > expires_at
    {
        return Err(AppError::Authentication("refresh token expired".into()));
    }

    // Look up current ACL role and contexts (propagates changes at refresh time)
    let (role, allowed_contexts) = check_acl_full(&state.acl_ks, &session.did).await?;

    // Generate new access token
    let config = state.config.read().await;
    let access_expiry = config.auth.access_token_expiry;
    drop(config);

    #[cfg(feature = "tee")]
    let tee_attested = state.tee.is_some();
    #[cfg(not(feature = "tee"))]
    let tee_attested = false;

    let claims = jwt_keys.new_claims(
        session.did.clone(),
        session.session_id.clone(),
        role.to_string(),
        allowed_contexts,
        access_expiry,
        tee_attested,
    );
    let access_expires_at = claims.exp;
    let access_token = jwt_keys.encode(&claims)?;

    info!(did = %session.did, session_id = %session.session_id, "token refreshed");
    audit!("auth.refresh", actor = &session.did, resource =&session.session_id, outcome = "success");

    Ok(Json(RefreshResponse {
        session_id: session.session_id,
        data: RefreshData {
            access_token,
            access_expires_at,
        },
    }))
}

// ---------- POST /auth/credentials ----------

#[derive(Debug, Deserialize)]
pub struct GenerateCredentialsRequest {
    pub role: Role,
    pub label: Option<String>,
    #[serde(default)]
    pub allowed_contexts: Vec<String>,
}

/// POST /auth/credentials — generate a new DID+keypair and register an ACL entry. Auth: Admin or Initiator.
pub async fn generate_credentials(
    auth: ManageAuth,
    State(state): State<AppState>,
    Json(req): Json<GenerateCredentialsRequest>,
) -> Result<(StatusCode, Json<GenerateCredentialsResultBody>), AppError> {
    let role_str = req.role.to_string();
    let result = operations::credentials::generate_credentials(
        &state.acl_ks,
        &state.config,
        &auth.0,
        req.role,
        req.label,
        req.allowed_contexts,
        "rest",
    )
    .await?;
    audit!("credentials.generate", actor = &auth.0.did, resource =&role_str, outcome = "success");
    Ok((StatusCode::CREATED, Json(result)))
}

// ---------- GET /auth/sessions ----------

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionSummary {
    pub session_id: String,
    pub did: String,
    pub state: SessionState,
    pub created_at: u64,
    pub refresh_expires_at: Option<u64>,
}

impl From<Session> for SessionSummary {
    fn from(s: Session) -> Self {
        Self {
            session_id: s.session_id,
            did: s.did,
            state: s.state,
            created_at: s.created_at,
            refresh_expires_at: s.refresh_expires_at,
        }
    }
}

/// GET /auth/sessions — list all active sessions. Auth: Admin or Initiator.
pub async fn session_list(
    _auth: ManageAuth,
    State(state): State<AppState>,
) -> Result<Json<Vec<SessionSummary>>, AppError> {
    let all = list_sessions(&state.sessions_ks).await?;
    let summaries: Vec<SessionSummary> = all.into_iter().map(SessionSummary::from).collect();
    info!(caller = %_auth.0.did, count = summaries.len(), "sessions listed");
    Ok(Json(summaries))
}

// ---------- DELETE /auth/sessions/{session_id} ----------

/// DELETE /auth/sessions/{session_id} — revoke a single session (own or admin). Auth: any authenticated user.
pub async fn revoke_session(
    auth: AuthClaims,
    State(state): State<AppState>,
    Path(session_id): Path<String>,
) -> Result<impl IntoResponse, AppError> {
    let session = get_session(&state.sessions_ks, &session_id)
        .await?
        .ok_or_else(|| AppError::NotFound(format!("session not found: {session_id}")))?;

    // Allow if caller owns the session or is admin
    if session.did != auth.did && auth.role != Role::Admin {
        return Err(AppError::Forbidden(
            "cannot revoke another user's session".into(),
        ));
    }

    delete_session(&state.sessions_ks, &session_id).await?;
    info!(caller = %auth.did, session_id = %session_id, "session revoked");
    audit!("session.revoke", actor = &auth.did, resource =&session_id, outcome = "success");
    Ok(StatusCode::NO_CONTENT)
}

// ---------- DELETE /auth/sessions?did=X ----------

#[derive(Debug, Deserialize)]
pub struct RevokeByDidQuery {
    pub did: String,
}

#[derive(Debug, Serialize)]
pub struct RevokeByDidResponse {
    pub revoked: u64,
}

/// DELETE /auth/sessions?did=X — revoke all sessions for a given DID. Auth: Admin only.
pub async fn revoke_sessions_by_did(
    _auth: AdminAuth,
    State(state): State<AppState>,
    Query(query): Query<RevokeByDidQuery>,
) -> Result<Json<RevokeByDidResponse>, AppError> {
    let all = list_sessions(&state.sessions_ks).await?;
    let mut revoked = 0u64;

    for session in all {
        if session.did == query.did {
            delete_session(&state.sessions_ks, &session.session_id).await?;
            revoked += 1;
        }
    }

    info!(caller = %_auth.0.did, target_did = %query.did, revoked, "sessions revoked by DID");
    audit!("session.revoke_by_did", actor = &_auth.0.did, resource =&query.did, outcome = "success");
    Ok(Json(RevokeByDidResponse { revoked }))
}