//! JWKS (JSON Web Key Set) — RFC 7517 + RFC 8037 OKP/Ed25519 publication.
//!
//! Phase 6/8 (RFC_2026-05-04_jwt-full-adoption §6.7 + §6.9). PAS publishes
//! its trusted Ed25519 verification keys at `/.well-known/jwks.json`;
//! consumers (chat-auth, pas-external SDK) fetch + cache to populate
//! `KeySet`.
//!
//! ── Pure RFC 7517 design (no extension fields) ─────────────────────────
//!
//! No `status`, no `cache_ttl_seconds`. The rotation lifecycle is "key in
//! JWKS = trusted; key removed = revoked", and TTL is communicated via
//! the `Cache-Control: max-age=N` HTTP header. The result is consumable
//! by any RFC 7517 library (`jsonwebtoken::jwk::JwkSet`, jose-jwk,
//! python-jose, ...) — no ppoppo-specific knowledge required.
//!
//! ── Shape (RFC 8037 §2 for Ed25519) ─────────────────────────────────────
//!
//! ```json
//! {
//!   "keys": [
//!     {"kty":"OKP","crv":"Ed25519","use":"sig","alg":"EdDSA","kid":"...","x":"<b64url(32B pubkey)>"}
//!   ]
//! }
//! ```
//!
//! `kty=OKP` (Octet Key Pair, RFC 8037), `crv=Ed25519`, `alg=EdDSA`. The
//! 32-byte public key is base64url-encoded without padding into `x`.
//! `use=sig` (signature; RFC 7517 §4.2).

use base64::Engine;
use jsonwebtoken::DecodingKey;
use serde::{Deserialize, Serialize};

use crate::KeySet;

/// ASN.1 DER prefix for an Ed25519 SubjectPublicKeyInfo (RFC 8410 §4).
/// Prepended to the raw 32-byte public key to form a 44-byte SPKI DER
/// blob that `jsonwebtoken::DecodingKey::from_ed_der` consumes directly.
///
/// Bytes in plain English:
/// - `30 2a` SEQUENCE, length 42
/// - `30 05` SEQUENCE, length 5 (AlgorithmIdentifier)
/// - `06 03 2b 65 70` OID 1.3.101.112 = id-Ed25519
/// - `03 21 00` BIT STRING length 33, 0 unused bits
const ED25519_SPKI_PREFIX: [u8; 12] =
    [0x30, 0x2a, 0x30, 0x05, 0x06, 0x03, 0x2b, 0x65, 0x70, 0x03, 0x21, 0x00];

/// JSON Web Key Set — collection of trusted public keys per RFC 7517 §5.
///
/// Equality + Clone derive enables ergonomic Arc-wrapping at the wiring
/// site without polluting the public surface.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Jwks {
    pub keys: Vec<Jwk>,
}

impl Jwks {
    /// Build a JWKS from a slice of (kid, 32-byte Ed25519 public key)
    /// tuples. Keys land in the order supplied — callers control which
    /// keys appear (typically: filter Revoked at the call site, supply
    /// Active + Retiring to the builder).
    #[must_use]
    pub fn from_ed25519_keys(keys: &[(&str, &[u8; 32])]) -> Self {
        Self {
            keys: keys
                .iter()
                .map(|(kid, pk)| Jwk::ed25519(kid, *pk))
                .collect(),
        }
    }

    /// Find the key with the matching kid that satisfies `use=sig`.
    /// Returns the 32-byte Ed25519 public key bytes when present and
    /// well-formed; `None` for missing kid or wrong key type. Used by
    /// consumer-side verification flows to bind a token's `kid` header
    /// to a trusted public key.
    #[must_use]
    pub fn find_ed25519(&self, kid: &str) -> Option<[u8; 32]> {
        let jwk = self.keys.iter().find(|k| k.kid == kid)?;
        jwk.ed25519_bytes()
    }

    /// Convert the JWKS into the engine's `KeySet`. Every well-formed
    /// `kty=OKP / crv=Ed25519` entry becomes a `(kid, DecodingKey)`
    /// binding; entries with any other shape are silently skipped (the
    /// engine cannot verify them anyway, and a future JWKS may legitimately
    /// carry mixed key types — RSA for legacy clients, EC for some federated
    /// IdP). The skip-or-fail tradeoff favours skip: a single malformed
    /// entry must not break key rotation for the well-formed siblings.
    ///
    /// Returns `Err(JwksError::DuplicateKid)` only if two entries share a
    /// kid — that is a control-plane bug (every kid is supposed to be
    /// globally unique), and admitting both would create non-determinism
    /// in `KeySet::get`.
    pub fn into_key_set(self) -> Result<KeySet, JwksError> {
        let mut key_set = KeySet::new();
        let mut seen: std::collections::HashSet<String> = Default::default();
        for jwk in self.keys {
            let Some(pk_bytes) = jwk.ed25519_bytes() else {
                continue;
            };
            if !seen.insert(jwk.kid.clone()) {
                return Err(JwksError::DuplicateKid(jwk.kid));
            }
            let mut der = Vec::with_capacity(ED25519_SPKI_PREFIX.len() + pk_bytes.len());
            der.extend_from_slice(&ED25519_SPKI_PREFIX);
            der.extend_from_slice(&pk_bytes);
            key_set.insert(jwk.kid, DecodingKey::from_ed_der(&der));
        }
        Ok(key_set)
    }
}

/// JWKS-side errors surfaced to consumers of `into_key_set`.
///
/// Distinct from `AuthError` because this fires at *configuration* time
/// (boot / cache refresh), not per-request verify time. Operators see
/// these in startup logs; users never do.
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum JwksError {
    /// Two JWK entries share a kid. Engine refuses to insert both
    /// because `KeySet::get` would be non-deterministic. Operator must
    /// fix the upstream JWKS source.
    #[error("duplicate kid in JWKS: '{0}'")]
    DuplicateKid(String),
}

/// A single JWK entry. Pinned to the OKP/Ed25519/EdDSA shape — other
/// `kty` values (`EC`, `RSA`, `oct`) deserialize but `ed25519_bytes()`
/// returns `None` so the engine never accidentally accepts a non-Ed25519
/// key.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Jwk {
    pub kty: String,
    pub crv: String,
    #[serde(rename = "use", default)]
    pub use_: String,
    pub alg: String,
    pub kid: String,
    pub x: String,
}

impl Jwk {
    /// Construct an Ed25519 JWK from its raw 32-byte public key. The
    /// fixed-string fields (`kty=OKP`, `crv=Ed25519`, `use=sig`,
    /// `alg=EdDSA`) match RFC 8037 §2 verbatim — every Ed25519 JWK PAS
    /// publishes carries this shape.
    #[must_use]
    pub fn ed25519(kid: &str, public_key: &[u8; 32]) -> Self {
        Self {
            kty: "OKP".to_string(),
            crv: "Ed25519".to_string(),
            use_: "sig".to_string(),
            alg: "EdDSA".to_string(),
            kid: kid.to_string(),
            x: base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(public_key),
        }
    }

    /// Decode the 32-byte Ed25519 public key carried in `x` when this
    /// JWK is shaped as `kty=OKP / crv=Ed25519`. Returns `None` for any
    /// other shape so consumers cannot accidentally feed a `kty=EC` or
    /// `kty=RSA` key to an Ed25519 verifier.
    #[must_use]
    pub fn ed25519_bytes(&self) -> Option<[u8; 32]> {
        if self.kty != "OKP" || self.crv != "Ed25519" {
            return None;
        }
        let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
            .decode(self.x.as_bytes())
            .ok()?;
        decoded.try_into().ok()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;

    fn sample_pubkey() -> [u8; 32] {
        // Deterministic test vector — first 32 bytes of a known Ed25519
        // public key. Value is irrelevant to the codec test; we only
        // care about round-trip fidelity.
        let mut bytes = [0u8; 32];
        for (i, b) in bytes.iter_mut().enumerate() {
            *b = i as u8;
        }
        bytes
    }

    #[test]
    fn ed25519_jwk_carries_rfc_8037_shape() {
        let jwk = Jwk::ed25519("k4.pid.test", &sample_pubkey());
        assert_eq!(jwk.kty, "OKP", "RFC 8037 §2: kty MUST be OKP for Ed25519");
        assert_eq!(jwk.crv, "Ed25519");
        assert_eq!(jwk.use_, "sig", "RFC 7517 §4.2: signature key");
        assert_eq!(jwk.alg, "EdDSA", "RFC 8037 §3.1: alg = EdDSA");
        assert_eq!(jwk.kid, "k4.pid.test");
    }

    #[test]
    fn ed25519_x_round_trips_through_base64url() {
        let pk = sample_pubkey();
        let jwk = Jwk::ed25519("kid-1", &pk);
        let recovered = jwk.ed25519_bytes().expect("must decode");
        assert_eq!(recovered, pk, "x must round-trip the raw public key bytes");
    }

    #[test]
    fn ed25519_x_is_base64url_no_pad() {
        // RFC 7515 §2 + RFC 7518: JWK fields use base64url without
        // padding. A `=` in the encoded form would be a wire-shape bug.
        let jwk = Jwk::ed25519("k", &sample_pubkey());
        assert!(
            !jwk.x.contains('='),
            "base64url MUST NOT carry padding: {}",
            jwk.x
        );
        assert!(
            !jwk.x.contains('+') && !jwk.x.contains('/'),
            "base64url MUST NOT use std-b64 chars: {}",
            jwk.x
        );
    }

    #[test]
    fn non_ed25519_kty_returns_none_from_bytes() {
        // Belt-and-suspenders for the engine: even if a JWKS publishes
        // an EC key with `x` carrying 32 bytes, ed25519_bytes refuses.
        let mut jwk = Jwk::ed25519("kid", &sample_pubkey());
        jwk.kty = "EC".to_string();
        assert!(jwk.ed25519_bytes().is_none(), "non-OKP must return None");
    }

    #[test]
    fn non_ed25519_crv_returns_none_from_bytes() {
        let mut jwk = Jwk::ed25519("kid", &sample_pubkey());
        jwk.crv = "X25519".to_string(); // OKP-but-key-agreement (RFC 8037 §3.2)
        assert!(
            jwk.ed25519_bytes().is_none(),
            "X25519 is OKP but for ECDH, not signing — must return None",
        );
    }

    #[test]
    fn jwks_round_trips_through_json() {
        let original = Jwks::from_ed25519_keys(&[
            ("kid-a", &sample_pubkey()),
            ("kid-b", &sample_pubkey()),
        ]);
        let json = serde_json::to_string(&original).unwrap();
        let parsed: Jwks = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, original, "JWKS must serde round-trip");
    }

    #[test]
    fn jwks_find_returns_matching_key() {
        let pk = sample_pubkey();
        let jwks = Jwks::from_ed25519_keys(&[("active-kid", &pk)]);
        let found = jwks
            .find_ed25519("active-kid")
            .expect("active-kid must be findable");
        assert_eq!(found, pk);
    }

    #[test]
    fn jwks_find_returns_none_for_unknown_kid() {
        let jwks = Jwks::from_ed25519_keys(&[("only-kid", &sample_pubkey())]);
        assert!(jwks.find_ed25519("missing-kid").is_none());
    }

    #[test]
    fn into_key_set_admits_well_formed_ed25519_entries() {
        let jwks = Jwks::from_ed25519_keys(&[
            ("kid-a", &sample_pubkey()),
            ("kid-b", &sample_pubkey()),
        ]);
        let key_set = jwks.into_key_set().expect("well-formed JWKS must convert");
        // KeySet::get is pub(crate); we don't have direct visibility into
        // its contents from here, but the absence of an error and the
        // duplicate-kid guard fires only when entries land — so success
        // here implies both entries inserted.
        let _ = key_set;
    }

    #[test]
    fn into_key_set_skips_non_ed25519_entries() {
        // A JWKS legitimately may carry other key types in a federation
        // scenario. ed25519_bytes returns None for them, so they get
        // silently skipped. Test by hand-constructing an EC-shaped entry
        // alongside a valid Ed25519 entry.
        let pk = sample_pubkey();
        let mut jwks = Jwks {
            keys: vec![
                Jwk::ed25519("ed-kid", &pk),
                Jwk {
                    kty: "EC".to_string(),
                    crv: "P-256".to_string(),
                    use_: "sig".to_string(),
                    alg: "ES256".to_string(),
                    kid: "ec-kid".to_string(),
                    x: "irrelevant".to_string(),
                },
            ],
        };
        // Sanity: the EC entry would fail Ed25519 decode.
        assert!(jwks.keys[1].ed25519_bytes().is_none());
        // Conversion succeeds — only the Ed25519 entry lands in KeySet.
        // (Non-Ed25519 silently skipped, no error.)
        let _ = jwks.into_key_set().expect("mixed-type JWKS must convert");
        // Add a duplicate to prove the dup-kid path is reachable.
        jwks = Jwks::from_ed25519_keys(&[("dup", &pk), ("dup", &pk)]);
        // KeySet doesn't impl Debug/PartialEq (it carries opaque
        // DecodingKey values from jsonwebtoken), so we assert on the
        // Err variant directly instead of through Result equality.
        let err = jwks
            .into_key_set()
            .err()
            .expect("duplicate kid must surface as Err");
        assert_eq!(err, JwksError::DuplicateKid("dup".to_string()));
    }

    #[test]
    fn into_key_set_round_trips_through_jwks_json_for_engine_verify() {
        // End-to-end smoke: a real signing key's public half goes into a
        // Jwks document, then comes back out as a KeySet that the engine
        // can use to verify a token signed by the matching private half.
        // This is the integration that 8.3 / 6.4 rely on.
        use crate::SigningKey;
        let (signer, _direct_key_set) = SigningKey::test_pair();

        // Reach into the test_pair PEM constants and re-derive the public
        // key bytes for the JWKS path. We use the documented test_pair
        // public-key PEM (from signing_key.rs) — base64 of the SPKI DER's
        // last 32 bytes.
        const TEST_PUBLIC_KEY_DER_B64: &str = "MCowBQYDK2VwAyEAh//e6j3It3xhjghg8Kpn2pM0jMCH/cvemGu4vv7D1Q4=";
        use base64::Engine as _;
        let der = base64::engine::general_purpose::STANDARD
            .decode(TEST_PUBLIC_KEY_DER_B64)
            .unwrap();
        // Last 32 bytes of the 44-byte SPKI are the raw public key.
        let pk_bytes: [u8; 32] = der[12..].try_into().unwrap();

        let jwks = Jwks::from_ed25519_keys(&[(signer.kid(), &pk_bytes)]);
        let _key_set = jwks.into_key_set().expect("well-formed JWKS must convert");
        // A full sign-then-verify round-trip lives in tests/keyset_jwks.rs;
        // here we only verify the conversion path itself.
    }

    #[test]
    fn jwks_json_shape_is_rfc_7517_compliant() {
        // Snapshot test — locks in the wire shape so a future refactor
        // (e.g. swapping to a different serde struct) cannot silently
        // drift to a non-standard shape.
        let pk = sample_pubkey();
        let jwks = Jwks::from_ed25519_keys(&[("test-kid", &pk)]);
        let value: serde_json::Value = serde_json::to_value(&jwks).unwrap();
        let key = &value["keys"][0];
        assert_eq!(key["kty"], "OKP");
        assert_eq!(key["crv"], "Ed25519");
        assert_eq!(key["use"], "sig");
        assert_eq!(key["alg"], "EdDSA");
        assert_eq!(key["kid"], "test-kid");
        assert!(key["x"].is_string());
        // No status, no cache_ttl_seconds — those are out (deliberate).
        assert!(value.get("cache_ttl_seconds").is_none());
        assert!(key.get("status").is_none());
    }
}
