//! End-to-end SAML SP tests.
//!
//! We mint a throwaway self-signed cert/key with `rcgen`, sign a SAML Response
//! template with `bergshamra::sign` (the same engine that verifies it), and
//! drive [`response::consume`] through the accept path and every reject path.
//! This proves both the signature wiring and the SP-semantic checks.

use time::format_description::well_known::Rfc3339;
use time::{Duration, OffsetDateTime};

use super::metadata::IdpMetadata;
use super::{response, SamlError, SpParams};

const SP_ENTITY: &str = "https://pxe.example.com";
const ACS: &str = "https://pxe.example.com/api/sso/acs";
const IDP_ENTITY: &str = "https://idp.example.com/realms/fleet";
const EMAIL: &str = "user@example.com";

struct TestIdp {
    cert_der: Vec<u8>,
    key_pem: String,
}

fn test_idp() -> TestIdp {
    let ck = rcgen::generate_simple_self_signed(vec!["idp.example.com".to_string()]).unwrap();
    TestIdp {
        cert_der: ck.cert.der().as_ref().to_vec(),
        key_pem: ck.key_pair.serialize_pem(),
    }
}

fn fmt(t: OffsetDateTime) -> String {
    t.replace_nanosecond(0).unwrap().format(&Rfc3339).unwrap()
}

/// Knobs for building a response template — defaults are a valid response.
struct Resp {
    issuer: String,
    audience: String,
    status: String,
    not_before: OffsetDateTime,
    not_on_or_after: OffsetDateTime,
    in_response_to: Option<String>,
    recipient: String,
}

impl Default for Resp {
    fn default() -> Self {
        let now = OffsetDateTime::now_utc();
        Self {
            issuer: IDP_ENTITY.into(),
            audience: SP_ENTITY.into(),
            status: "urn:oasis:names:tc:SAML:2.0:status:Success".into(),
            not_before: now - Duration::minutes(5),
            not_on_or_after: now + Duration::hours(1),
            in_response_to: Some("_req-abc".into()),
            recipient: ACS.into(),
        }
    }
}

impl Resp {
    /// The unsigned template (a `<ds:Signature>` with empty values).
    fn template(&self) -> String {
        let now = fmt(OffsetDateTime::now_utc());
        let irt = self
            .in_response_to
            .as_ref()
            .map(|v| format!(r#" InResponseTo="{v}""#))
            .unwrap_or_default();
        format!(
            r##"<samlp:Response xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol" xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion" ID="_resp1" Version="2.0" IssueInstant="{now}" Destination="{ACS}"{irt}>
  <saml:Issuer>{issuer}</saml:Issuer>
  <samlp:Status><samlp:StatusCode Value="{status}"/></samlp:Status>
  <saml:Assertion ID="_assertion1" Version="2.0" IssueInstant="{now}">
    <saml:Issuer>{issuer}</saml:Issuer>
    <ds:Signature xmlns:ds="http://www.w3.org/2000/09/xmldsig#">
      <ds:SignedInfo>
        <ds:CanonicalizationMethod Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
        <ds:SignatureMethod Algorithm="http://www.w3.org/2001/04/xmldsig-more#ecdsa-sha256"/>
        <ds:Reference URI="#_assertion1">
          <ds:Transforms>
            <ds:Transform Algorithm="http://www.w3.org/2000/09/xmldsig#enveloped-signature"/>
            <ds:Transform Algorithm="http://www.w3.org/2001/10/xml-exc-c14n#"/>
          </ds:Transforms>
          <ds:DigestMethod Algorithm="http://www.w3.org/2001/04/xmlenc#sha256"/>
          <ds:DigestValue></ds:DigestValue>
        </ds:Reference>
      </ds:SignedInfo>
      <ds:SignatureValue></ds:SignatureValue>
    </ds:Signature>
    <saml:Subject>
      <saml:NameID Format="urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress">{EMAIL}</saml:NameID>
      <saml:SubjectConfirmation Method="urn:oasis:names:tc:SAML:2.0:cm:bearer">
        <saml:SubjectConfirmationData Recipient="{recipient}" NotOnOrAfter="{noa}"{irt}/>
      </saml:SubjectConfirmation>
    </saml:Subject>
    <saml:Conditions NotBefore="{nb}" NotOnOrAfter="{noa}">
      <saml:AudienceRestriction><saml:Audience>{audience}</saml:Audience></saml:AudienceRestriction>
    </saml:Conditions>
    <saml:AuthnStatement AuthnInstant="{now}" SessionIndex="sess-123">
      <saml:AuthnContext><saml:AuthnContextClassRef>urn:oasis:names:tc:SAML:2.0:ac:classes:Password</saml:AuthnContextClassRef></saml:AuthnContext>
    </saml:AuthnStatement>
    <saml:AttributeStatement>
      <saml:Attribute Name="displayName"><saml:AttributeValue>Miles Ward</saml:AttributeValue></saml:Attribute>
    </saml:AttributeStatement>
  </saml:Assertion>
</samlp:Response>"##,
            issuer = self.issuer,
            status = self.status,
            audience = self.audience,
            recipient = self.recipient,
            nb = fmt(self.not_before),
            noa = fmt(self.not_on_or_after),
        )
    }
}

fn sign(template: &str, key_pem: &str) -> String {
    let key = bergshamra::keys::loader::load_pem_auto(key_pem.as_bytes(), None)
        .expect("load test signing key");
    let mut km = bergshamra::keys::KeysManager::new();
    km.add_key(key);
    let ctx = bergshamra::DsigContext::new(km);
    bergshamra::sign(&ctx, template).expect("sign test response")
}

fn sp() -> SpParams {
    SpParams {
        entity_id: SP_ENTITY.into(),
        acs_url: ACS.into(),
    }
}

fn idp(cert_der: Vec<u8>) -> IdpMetadata {
    IdpMetadata {
        entity_id: IDP_ENTITY.into(),
        sso_redirect_url: None,
        sso_post_url: None,
        signing_certs_der: vec![cert_der],
    }
}

fn consume(xml: &str, cert_der: Vec<u8>) -> Result<response::VerifiedResponse, SamlError> {
    response::consume(
        xml,
        &sp(),
        &idp(cert_der),
        OffsetDateTime::now_utc(),
        Duration::seconds(60),
    )
}

#[test]
fn good_response_yields_principal() {
    let t = test_idp();
    let signed = sign(&Resp::default().template(), &t.key_pem);
    let out = consume(&signed, t.cert_der).expect("valid response should verify");
    assert_eq!(out.principal.email, EMAIL);
    assert_eq!(out.principal.name_id, EMAIL);
    assert_eq!(out.principal.display_name.as_deref(), Some("Miles Ward"));
    assert_eq!(out.in_response_to.as_deref(), Some("_req-abc"));
    assert_eq!(out.assertion_id, "_assertion1");
    assert_eq!(out.session_index.as_deref(), Some("sess-123"));
}

#[test]
fn tampered_assertion_is_rejected() {
    let t = test_idp();
    let signed = sign(&Resp::default().template(), &t.key_pem);
    // Flip the subject email after signing — breaks the digest.
    let tampered = signed.replace(EMAIL, "attacker@evil.example");
    assert_ne!(signed, tampered);
    assert!(matches!(
        consume(&tampered, t.cert_der),
        Err(SamlError::Signature(_) | SamlError::SignatureScope)
    ));
}

#[test]
fn unsigned_response_is_rejected() {
    let t = test_idp();
    // Feed the *unsigned* template (empty SignatureValue) straight in.
    let unsigned = Resp::default().template();
    assert!(matches!(
        consume(&unsigned, t.cert_der),
        Err(SamlError::Signature(_))
    ));
}

#[test]
fn wrong_signing_key_is_rejected() {
    let signer = test_idp();
    let other = test_idp(); // different keypair pinned as the "IdP" cert
    let signed = sign(&Resp::default().template(), &signer.key_pem);
    assert!(matches!(
        consume(&signed, other.cert_der),
        Err(SamlError::Signature(_))
    ));
}

#[test]
fn wrong_audience_is_rejected() {
    let t = test_idp();
    let r = Resp {
        audience: "https://someone-else.example".into(),
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    assert!(matches!(
        consume(&signed, t.cert_der),
        Err(SamlError::AudienceMismatch)
    ));
}

#[test]
fn expired_assertion_is_rejected() {
    let t = test_idp();
    let now = OffsetDateTime::now_utc();
    let r = Resp {
        not_before: now - Duration::hours(2),
        not_on_or_after: now - Duration::hours(1),
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    assert!(matches!(
        consume(&signed, t.cert_der),
        Err(SamlError::TimeBounds)
    ));
}

#[test]
fn future_assertion_is_rejected() {
    let t = test_idp();
    let now = OffsetDateTime::now_utc();
    let r = Resp {
        not_before: now + Duration::hours(1),
        not_on_or_after: now + Duration::hours(2),
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    assert!(matches!(
        consume(&signed, t.cert_der),
        Err(SamlError::TimeBounds)
    ));
}

#[test]
fn wrong_issuer_is_rejected() {
    let t = test_idp();
    let r = Resp {
        issuer: "https://evil-idp.example".into(),
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    assert!(matches!(
        consume(&signed, t.cert_der),
        Err(SamlError::IssuerMismatch)
    ));
}

#[test]
fn non_success_status_is_rejected() {
    let t = test_idp();
    let r = Resp {
        status: "urn:oasis:names:tc:SAML:2.0:status:Requester".into(),
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    assert!(matches!(
        consume(&signed, t.cert_der),
        Err(SamlError::Status(_))
    ));
}

#[test]
fn idp_initiated_has_no_in_response_to() {
    // No InResponseTo => the HTTP layer must gate it behind allow_idp_initiated.
    let t = test_idp();
    let r = Resp {
        in_response_to: None,
        ..Resp::default()
    };
    let signed = sign(&r.template(), &t.key_pem);
    let out = consume(&signed, t.cert_der).expect("unsolicited but otherwise valid");
    assert!(out.in_response_to.is_none());
}
