authia 0.3.4

use crate::error::AuthiaError;
use crate::types::{TokenPayload, VerifyOptions};
use base64::prelude::*;
use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use serde::de::DeserializeOwned;
use serde_json::Value;
use std::sync::OnceLock;
use subtle::ConstantTimeEq;

// クロックスキュー許容時間（5分）
const CLOCK_SKEW_SECONDS: i64 = 300;

// 公開鍵キャッシュ（オプション機能）
// Note: 単一の公開鍵のみを使用する場合に有効
// 公開鍵が変更される場合は、プロセスの再起動が必要
static CACHED_PUBLIC_KEY: OnceLock<(String, VerifyingKey)> = OnceLock::new();

pub fn verify_jwt<T: DeserializeOwned + TokenPayload>(
    token: &str,
    options: &VerifyOptions,
) -> Result<T, AuthiaError> {
    // Split JWT into parts
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(AuthiaError::malformed_token("JWT must have 3 parts"));
    }

    let header_b64 = parts[0];
    let payload_b64 = parts[1];
    let signature_b64 = parts[2];

    // Verify algorithm (fast path: common header for EdDSA)
    // Most EdDSA headers are {"alg":"EdDSA","typ":"JWT"} or Similar.
    // We only decode if we really need to, but here for security we decode once.
    let header_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(header_b64)
        .map_err(|_| AuthiaError::malformed_token("Invalid header encoding"))?;
    
    if !header_bytes.windows(5).any(|w| w == b"EdDSA") {
        return Err(AuthiaError::invalid_signature());
    }

    // Decode payload
    let payload_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(payload_b64)
        .map_err(|e| AuthiaError::malformed_token(format!("Invalid payload encoding: {}", e)))?;

    // DESERIALIZE ONCE (Optimized: Parse directly into final struct)
    let payload: T = serde_json::from_slice(&payload_bytes)
        .map_err(|e| AuthiaError::invalid_claims(format!("Failed to parse payload: {}", e)))?;

    // Perform validation using traits (faster than Value access)
    let now = js_sys::Date::now() as i64 / 1000;
    payload.validate(options, now)?;

    // Decode signature
    let signature_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(signature_b64)
        .map_err(|e| AuthiaError::malformed_token(format!("Invalid signature encoding: {}", e)))?;

    let signature =
        Signature::from_slice(&signature_bytes).map_err(|_| AuthiaError::invalid_signature())?;

    // Get or decode public key (with caching)
    let public_key = get_or_decode_public_key(options)?;

    // Verify signature (optimized message building)
    let mut message = Vec::with_capacity(header_b64.len() + 1 + payload_b64.len());
    message.extend_from_slice(header_b64.as_bytes());
    message.push(b'.');
    message.extend_from_slice(payload_b64.as_bytes());

    public_key
        .verify(&message, &signature)
        .map_err(|_| AuthiaError::invalid_signature())?;

    Ok(payload)
}

/// Get public key from cache or decode it
fn get_or_decode_public_key(options: &VerifyOptions) -> Result<&'static VerifyingKey, AuthiaError> {
    // Priority 1: Try raw JWK (faster, no base64 decode)
    if let Some(raw_jwk) = &options.public_key_jwk_raw {
        // Try to get from cache
        if let Some((cached_jwk, cached_key)) = CACHED_PUBLIC_KEY.get() {
            if cached_jwk == raw_jwk {
                return Ok(cached_key);
            }
        }

        // Decode from raw JSON and cache
        let key = decode_public_key_from_raw_json(raw_jwk)?;

        // Try to initialize cache (may fail if already initialized with different key)
        let _ = CACHED_PUBLIC_KEY.set((raw_jwk.to_string(), key));

        // Get from cache (guaranteed to exist now)
        return CACHED_PUBLIC_KEY
            .get()
            .map(|(_, k)| k)
            .ok_or_else(|| AuthiaError::invalid_public_key("Failed to cache public key"));
    }

    // Priority 2: Try base64-encoded JWK (legacy, backward compatibility)
    if let Some(jwk_base64) = &options.public_key_jwk {
        // Try to get from cache
        if let Some((cached_jwk, cached_key)) = CACHED_PUBLIC_KEY.get() {
            if cached_jwk == jwk_base64 {
                return Ok(cached_key);
            }
        }

        // Decode from base64 and cache
        let key = decode_public_key_jwk(jwk_base64)?;

        // Try to initialize cache (may fail if already initialized with different key)
        let _ = CACHED_PUBLIC_KEY.set((jwk_base64.to_string(), key));

        // Get from cache (guaranteed to exist now)
        return CACHED_PUBLIC_KEY
            .get()
            .map(|(_, k)| k)
            .ok_or_else(|| AuthiaError::invalid_public_key("Failed to cache public key"));
    }

    Err(AuthiaError::invalid_public_key(
        "No public key provided (either publicKeyJwk or publicKeyJwkRaw required)",
    ))
}

/// Decode public key from raw JWK JSON string (auto-detects if it's base64 encoded)
fn decode_public_key_from_raw_json(jwk_input: &str) -> Result<VerifyingKey, AuthiaError> {
    let jwk_trimmed = jwk_input.trim();
    
    // Auto-detect: if it doesn't start with '{', it's likely base64 encoded (legacy environment)
    let jwk_json_bytes = if jwk_trimmed.starts_with('{') {
        jwk_trimmed.as_bytes().to_vec()
    } else {
        BASE64_STANDARD
            .decode(jwk_trimmed)
            .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid public key format: expected JSON but could not decode as Base64: {}", e)))?
    };

    let jwk: Value = serde_json::from_slice(&jwk_json_bytes)
        .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid JWK JSON: {}", e)))?;

    // Verify kty parameter
    let kty = jwk
        .get("kty")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'kty' parameter in JWK"))?;

    if kty != "OKP" {
        return Err(AuthiaError::invalid_public_key(format!(
            "Invalid 'kty': expected 'OKP', got '{}'",
            kty
        )));
    }

    // Verify crv parameter
    let crv = jwk
        .get("crv")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'crv' parameter in JWK"))?;

    if crv != "Ed25519" {
        return Err(AuthiaError::invalid_public_key(format!(
            "Invalid 'crv': expected 'Ed25519', got '{}'",
            crv
        )));
    }

    // Extract 'x' parameter (public key)
    let x = jwk
        .get("x")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'x' parameter in JWK"))?;

    // Decode the public key bytes
    let key_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(x)
        .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid 'x' encoding: {}", e)))?;

    // Create VerifyingKey
    VerifyingKey::from_bytes(
        key_bytes.as_slice().try_into().map_err(|_| {
            AuthiaError::invalid_public_key("Invalid key length (expected 32 bytes)")
        })?,
    )
    .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid Ed25519 key: {}", e)))
}

/// Decode public key from base64-encoded JWK (legacy)
fn decode_public_key_jwk(jwk_base64: &str) -> Result<VerifyingKey, AuthiaError> {
    // Decode base64-encoded JWK
    let jwk_json = BASE64_STANDARD
        .decode(jwk_base64)
        .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid base64: {}", e)))?;

    let jwk: Value = serde_json::from_slice(&jwk_json)
        .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid JWK JSON: {}", e)))?;

    // Verify kty parameter
    let kty = jwk
        .get("kty")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'kty' parameter in JWK"))?;

    if kty != "OKP" {
        return Err(AuthiaError::invalid_public_key(format!(
            "Invalid 'kty': expected 'OKP', got '{}'",
            kty
        )));
    }

    // Verify crv parameter
    let crv = jwk
        .get("crv")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'crv' parameter in JWK"))?;

    if crv != "Ed25519" {
        return Err(AuthiaError::invalid_public_key(format!(
            "Invalid 'crv': expected 'Ed25519', got '{}'",
            crv
        )));
    }

    // Extract 'x' parameter (public key)
    let x = jwk
        .get("x")
        .and_then(|v| v.as_str())
        .ok_or_else(|| AuthiaError::invalid_public_key("Missing 'x' parameter in JWK"))?;

    // Decode the public key bytes
    let key_bytes = BASE64_URL_SAFE_NO_PAD
        .decode(x)
        .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid 'x' encoding: {}", e)))?;

    // Create VerifyingKey
    VerifyingKey::from_bytes(
        key_bytes.as_slice().try_into().map_err(|_| {
            AuthiaError::invalid_public_key("Invalid key length (expected 32 bytes)")
        })?,
    )
    .map_err(|e| AuthiaError::invalid_public_key(format!("Invalid Ed25519 key: {}", e)))
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_jwt_split() {
        let token = "header.payload.signature";
        let parts: Vec<&str> = token.split('.').collect();
        assert_eq!(parts.len(), 3);
    }
}