//! Agent Card JWS signing.
//!
//! Produces RFC 7515 detached JSON Web Signatures over an
//! [`AgentCard`] so peers can verify provenance / domain.
//!
//! ## Canonicalization
//!
//! The signing payload is canonicalized via RFC 8785 (JSON
//! Canonicalization Scheme) before being base64url-encoded into the
//! JWS signing input. A verifier can re-serialize the `AgentCard`
//! JSON using *their own* JSON library, run the same JCS pass, and
//! reproduce byte-identical signing input — making
//! cross-implementation verification reliable. Without JCS, byte
//! differences between Rust's `serde_json` ordering and Python's
//! `json.dumps` ordering would have broken every cross-impl
//! verifier.

use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use jsonwebtoken::{crypto, Algorithm, DecodingKey, EncodingKey, Header};

use crate::types::{AgentCard, AgentCardSignature};

#[derive(Debug, thiserror::Error)]
pub enum SigningError {
    #[error("could not load signing key: {0}")]
    Key(#[from] jsonwebtoken::errors::Error),
    #[error("could not serialize agent card: {0}")]
    Serialize(#[from] serde_json::Error),
    #[error("internal signing error: {0}")]
    Internal(String),
    #[error("agent card has no signatures")]
    Unsigned,
    /// Returned when at least one attached signature reached
    /// `crypto::verify` and that call returned `Ok(false)` —
    /// genuine signature mismatch (wrong key, tampered card,
    /// truncated signature bytes).
    #[error("agent card signature failed verification")]
    BadSignature,
    /// Returned when *every* attached signature was unparseable
    /// before reaching `crypto::verify`: header wasn't valid
    /// base64url, header JSON didn't have an `alg`, or `alg` named
    /// an algorithm `jsonwebtoken` doesn't recognise. Distinct from
    /// `BadSignature` so security audits can tell "wrong key /
    /// tampering" from "buggy peer / wire corruption."
    #[error("every attached signature has an unparseable protected header")]
    MalformedSignatureHeader,
}

/// Bundle of inputs needed to sign an `AgentCard`.
pub struct CardSigner {
    encoding_key: EncodingKey,
    algorithm: Algorithm,
    key_id: Option<String>,
}

impl CardSigner {
    /// HMAC-SHA256 (HS256) signer. Useful for tests and for shared-
    /// secret deployments. Production deployments prefer the
    /// asymmetric variants below.
    pub fn from_hs256_secret(secret: &[u8]) -> Self {
        Self {
            encoding_key: EncodingKey::from_secret(secret),
            algorithm: Algorithm::HS256,
            key_id: None,
        }
    }

    /// RSA SHA-256 (RS256) signer. `pem` is the PKCS#1 or PKCS#8
    /// encoded private key.
    pub fn from_rs256_pem(pem: &[u8]) -> Result<Self, SigningError> {
        Ok(Self {
            encoding_key: EncodingKey::from_rsa_pem(pem)?,
            algorithm: Algorithm::RS256,
            key_id: None,
        })
    }

    /// ECDSA P-256 / SHA-256 (ES256) signer. `pem` is the PKCS#8
    /// encoded private key.
    pub fn from_es256_pem(pem: &[u8]) -> Result<Self, SigningError> {
        Ok(Self {
            encoding_key: EncodingKey::from_ec_pem(pem)?,
            algorithm: Algorithm::ES256,
            key_id: None,
        })
    }

    /// Attach a `kid` (key id) to the protected header. Peers use it
    /// to pick the right verification key from a JWKS.
    pub fn with_key_id(mut self, key_id: impl Into<String>) -> Self {
        self.key_id = Some(key_id.into());
        self
    }
}

/// Sign an `AgentCard` and append the signature to its `signatures`
/// vector. The card's pre-existing signatures are cleared from the
/// signing payload (signing the card without its own signatures is
/// the only correct order — otherwise signatures would chain).
///
/// Both the protected header and the payload are JCS-canonicalised
/// before being base64url-encoded into the JWS signing input.
pub fn sign_agent_card(card: &mut AgentCard, signer: &CardSigner) -> Result<(), SigningError> {
    let mut header = Header::new(signer.algorithm);
    header.kid = signer.key_id.clone();

    let mut payload_card = card.clone();
    payload_card.signatures.clear();

    // Header: serialised normally. Verifiers don't recompute the
    // header — they base64url-decode the on-card `protected` field
    // and read the JSON. Canonicalising here would harmlessly reorder
    // keys, but custom headers (e.g. `crit` + extension order, or
    // pre-computed RFC 7638 jwk thumbprints) could surprise peers
    // that compare bytes. So: serde_json for the header, JCS only
    // for the payload (which is the load-bearing canonicalisation).
    let header_bytes = serde_json::to_vec(&header)?;
    let payload_canon = serde_jcs::to_vec(&payload_card)?;
    let header_b64 = URL_SAFE_NO_PAD.encode(&header_bytes);
    let payload_b64 = URL_SAFE_NO_PAD.encode(&payload_canon);
    let signing_input = format!("{}.{}", header_b64, payload_b64);

    let signature_b64 = crypto::sign(
        signing_input.as_bytes(),
        &signer.encoding_key,
        signer.algorithm,
    )?;

    card.signatures.push(AgentCardSignature {
        protected: header_b64,
        signature: signature_b64,
        header: None,
    });
    Ok(())
}

/// Verify an `AgentCard`'s signature.
///
/// Returns the index of the first signature on `card.signatures`
/// that validates against `key`. Returns `SigningError::Unsigned`
/// when the card has no signatures, `BadSignature` when none of the
/// attached signatures validates.
///
/// The verifier reproduces the JCS-canonical signing input:
/// canonicalize the card with `signatures` cleared, base64url-encode,
/// concatenate with the on-card `protected` header. So long as the
/// embedder used [`sign_agent_card`] (or any signer that takes the
/// same JCS-canonical bytes), verification succeeds.
pub fn verify_agent_card(card: &AgentCard, key: &DecodingKey) -> Result<usize, SigningError> {
    if card.signatures.is_empty() {
        return Err(SigningError::Unsigned);
    }
    let mut payload_card = card.clone();
    payload_card.signatures.clear();
    let payload_canon = serde_jcs::to_vec(&payload_card)?;
    let payload_b64 = URL_SAFE_NO_PAD.encode(&payload_canon);

    let mut tried_verify = false;
    for (idx, sig) in card.signatures.iter().enumerate() {
        // Decode the protected header to learn the algorithm. A
        // malformed header on signature N must not stop us from
        // trying signature N+1 — peers may publish multiple
        // signatures (key rotation, dual-alg) and one being broken
        // shouldn't DoS the rest.
        let Ok(header_bytes) = URL_SAFE_NO_PAD.decode(&sig.protected) else {
            continue;
        };
        let Ok(alg) = serde_json::from_slice::<HeaderAlgOnly>(&header_bytes) else {
            continue;
        };
        // Parse the alg string through serde so any algorithm
        // `jsonwebtoken` supports (HS{256,384,512}, RS{256,384,512},
        // PS{256,384,512}, ES{256,384}, EdDSA) round-trips.
        let Ok(algorithm) = serde_json::from_value::<Algorithm>(serde_json::Value::String(alg.alg))
        else {
            continue;
        };

        let signing_input = format!("{}.{}", sig.protected, payload_b64);
        tried_verify = true;
        match crypto::verify(&sig.signature, signing_input.as_bytes(), key, algorithm) {
            Ok(true) => return Ok(idx),
            Ok(false) | Err(_) => continue,
        }
    }
    if tried_verify {
        Err(SigningError::BadSignature)
    } else {
        Err(SigningError::MalformedSignatureHeader)
    }
}

#[derive(serde::Deserialize)]
struct HeaderAlgOnly {
    alg: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{AgentCapabilities, AgentInterface, AgentProvider, TransportProtocol};
    use std::collections::HashMap;

    fn unsigned_card() -> AgentCard {
        AgentCard {
            name: "test".into(),
            description: "for signing".into(),
            url: "https://example.test".into(),
            version: "1.0.0".into(),
            protocol_version: "1.0".into(),
            preferred_transport: Some("JSONRPC".into()),
            provider: AgentProvider {
                organization: "Parslee".into(),
                url: None,
            },
            capabilities: AgentCapabilities::default(),
            default_input_modes: vec!["text".into()],
            default_output_modes: vec!["text".into()],
            skills: vec![],
            documentation_url: None,
            icon_url: None,
            supported_interfaces: vec![],
            additional_interfaces: vec![AgentInterface {
                url: "https://example.test".into(),
                protocol_binding: "JSONRPC".into(),
                transport: Some(TransportProtocol::JsonRpc),
                tenant: None,
                protocol_version: "1.0".into(),
            }],
            security_schemes: HashMap::new(),
            supports_authenticated_extended_card: false,
            security_requirements: vec![],
            signatures: vec![],
        }
    }

    #[test]
    fn hs256_signature_appended_with_kid() {
        let signer = CardSigner::from_hs256_secret(b"shared-secret").with_key_id("test-key");
        let mut card = unsigned_card();
        sign_agent_card(&mut card, &signer).expect("signed");
        assert_eq!(card.signatures.len(), 1);
        let sig = &card.signatures[0];
        assert!(!sig.protected.is_empty());
        assert!(!sig.signature.is_empty());

        // Decode the protected header and check kid.
        let header_bytes = URL_SAFE_NO_PAD.decode(&sig.protected).unwrap();
        let header_json: serde_json::Value = serde_json::from_slice(&header_bytes).unwrap();
        assert_eq!(header_json["kid"], "test-key");
        assert_eq!(header_json["alg"], "HS256");
    }

    #[test]
    fn signing_clears_pre_existing_signatures_before_payload() {
        let signer = CardSigner::from_hs256_secret(b"k");
        let mut card = unsigned_card();
        // Sign once.
        sign_agent_card(&mut card, &signer).expect("first sign");
        let first_sig = card.signatures[0].signature.clone();
        // Sign again — the second signing should treat the card as
        // if it had no signatures (i.e. payload bytes are identical
        // to the first signing). So both signatures should be byte-
        // identical for HS256/deterministic.
        sign_agent_card(&mut card, &signer).expect("second sign");
        assert_eq!(card.signatures.len(), 2);
        assert_eq!(card.signatures[1].signature, first_sig);
    }

    #[test]
    fn sign_then_verify_round_trips() {
        let secret = b"shared-secret-for-test";
        let signer = CardSigner::from_hs256_secret(secret).with_key_id("k1");
        let mut card = unsigned_card();
        sign_agent_card(&mut card, &signer).expect("sign");

        let key = DecodingKey::from_secret(secret);
        let idx = verify_agent_card(&card, &key).expect("verify");
        assert_eq!(idx, 0);
    }

    #[test]
    fn verify_after_json_round_trip_succeeds() {
        // The wire-form round-trip — peers receive the card as JSON,
        // parse it, and re-serialize on their side. JCS makes that
        // signing-input-byte-identical so verification still works.
        let secret = b"k";
        let signer = CardSigner::from_hs256_secret(secret);
        let mut card = unsigned_card();
        sign_agent_card(&mut card, &signer).expect("sign");

        let json = serde_json::to_string(&card).expect("ser");
        let reparsed: AgentCard = serde_json::from_str(&json).expect("deser");

        let key = DecodingKey::from_secret(secret);
        verify_agent_card(&reparsed, &key).expect("verify after round-trip");
    }

    #[test]
    fn verify_rejects_tampered_card() {
        let secret = b"k";
        let signer = CardSigner::from_hs256_secret(secret);
        let mut card = unsigned_card();
        sign_agent_card(&mut card, &signer).expect("sign");

        // Tamper: change the agent description.
        card.description = "TAMPERED".into();

        let key = DecodingKey::from_secret(secret);
        assert!(matches!(
            verify_agent_card(&card, &key),
            Err(SigningError::BadSignature)
        ));
    }

    #[test]
    fn verify_rejects_unsigned_card() {
        let card = unsigned_card();
        let key = DecodingKey::from_secret(b"k");
        assert!(matches!(
            verify_agent_card(&card, &key),
            Err(SigningError::Unsigned)
        ));
    }
}
