use crate::traits::Partition;

/// Partition with cached SK/IK key id strings.
///
/// `system_key_id`/`intermediate_key_id` are called on every encrypt
/// and decrypt to look up keys in the metastore and SK/IK caches. The
/// previous implementation rebuilt the strings via `format!` per call,
/// which allocates. Cache them at construction since neither id can
/// change for the lifetime of the partition. T-finding "SK id
/// allocated via format! per encrypt" in
/// `docs/review-2026-05-05-findings.md`.
#[derive(Clone, Debug)]
pub struct DefaultPartition {
    // The original (id, service, product, suffix) inputs are kept so
    // `Debug` formatting and any future accessors can recover them
    // without re-parsing the cached SK/IK strings.
    #[allow(dead_code)]
    id: String,
    #[allow(dead_code)]
    service: String,
    #[allow(dead_code)]
    product: String,
    #[allow(dead_code)]
    suffix: Option<String>,
    cached_sk_id: String,
    cached_ik_id: String,
    cached_ik_validation_prefix: Option<String>,
}

impl DefaultPartition {
    pub fn new(id: String, service: String, product: String) -> Self {
        let cached_sk_id = format!("_SK_{}_{}", service, product);
        let cached_ik_id = format!("_IK_{}_{}_{}", id, service, product);
        Self {
            id,
            service,
            product,
            suffix: None,
            cached_sk_id,
            cached_ik_id,
            cached_ik_validation_prefix: None,
        }
    }
    pub fn new_suffixed(id: String, service: String, product: String, suffix: String) -> Self {
        let cached_sk_id = format!("_SK_{}_{}_{}", service, product, suffix);
        let cached_ik_id = format!("_IK_{}_{}_{}_{}", id, service, product, suffix);
        let cached_ik_validation_prefix = Some(format!("_IK_{}_{}_{}_", id, service, product));
        Self {
            id,
            service,
            product,
            suffix: Some(suffix),
            cached_sk_id,
            cached_ik_id,
            cached_ik_validation_prefix,
        }
    }

    /// Returns the prefix used to validate IK ids when the partition has a suffix.
    /// Returns `None` when there is no suffix (exact match only).
    pub fn ik_validation_prefix(&self) -> Option<String> {
        self.cached_ik_validation_prefix.clone()
    }
}

impl Partition for DefaultPartition {
    fn system_key_id(&self) -> String {
        self.cached_sk_id.clone()
    }
    fn intermediate_key_id(&self) -> String {
        self.cached_ik_id.clone()
    }
    fn is_valid_intermediate_key_id(&self, id: &str) -> bool {
        match &self.cached_ik_validation_prefix {
            Some(prefix) => id == self.cached_ik_id || id.starts_with(prefix.as_str()),
            None => id == self.cached_ik_id,
        }
    }
}

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

    #[test]
    fn system_key_id_without_suffix() {
        let p = DefaultPartition::new("user1".into(), "svc".into(), "prod".into());
        assert_eq!(p.system_key_id(), "_SK_svc_prod");
    }

    #[test]
    fn system_key_id_with_suffix() {
        let p = DefaultPartition::new_suffixed(
            "user1".into(),
            "svc".into(),
            "prod".into(),
            "us-east-1".into(),
        );
        assert_eq!(p.system_key_id(), "_SK_svc_prod_us-east-1");
    }

    #[test]
    fn intermediate_key_id_without_suffix() {
        let p = DefaultPartition::new("user1".into(), "svc".into(), "prod".into());
        assert_eq!(p.intermediate_key_id(), "_IK_user1_svc_prod");
    }

    #[test]
    fn intermediate_key_id_with_suffix() {
        let p = DefaultPartition::new_suffixed(
            "user1".into(),
            "svc".into(),
            "prod".into(),
            "us-west-2".into(),
        );
        assert_eq!(p.intermediate_key_id(), "_IK_user1_svc_prod_us-west-2");
    }

    #[test]
    fn is_valid_ik_id_exact_match_no_suffix() {
        let p = DefaultPartition::new("u".into(), "s".into(), "p".into());
        assert!(p.is_valid_intermediate_key_id("_IK_u_s_p"));
        assert!(!p.is_valid_intermediate_key_id("_IK_u_s_p_extra"));
        assert!(!p.is_valid_intermediate_key_id("_IK_other_s_p"));
        assert!(!p.is_valid_intermediate_key_id(""));
    }

    #[test]
    fn is_valid_ik_id_with_suffix_accepts_prefix() {
        let p = DefaultPartition::new_suffixed("u".into(), "s".into(), "p".into(), "r1".into());
        // Exact match
        assert!(p.is_valid_intermediate_key_id("_IK_u_s_p_r1"));
        // Prefix match (different suffix)
        assert!(p.is_valid_intermediate_key_id("_IK_u_s_p_r2"));
        // Wrong partition id
        assert!(!p.is_valid_intermediate_key_id("_IK_other_s_p_r1"));
        // Must not accept non-suffixed IK (missing trailing delimiter)
        assert!(!p.is_valid_intermediate_key_id("_IK_u_s_p"));
    }

    #[test]
    fn different_partitions_produce_different_key_ids() {
        let p1 = DefaultPartition::new("a".into(), "svc".into(), "prod".into());
        let p2 = DefaultPartition::new("b".into(), "svc".into(), "prod".into());
        assert_eq!(p1.system_key_id(), p2.system_key_id());
        assert_ne!(p1.intermediate_key_id(), p2.intermediate_key_id());
    }
}
