//! Predicate alias storage for semantic predicate matching.
//!
//! Persists predicate alias sets so that semantic equivalences survive restarts.
//! When a Trust Pack defines `enabled ↔ required`, this store ensures the alias
//! is available on subsequent scans.
//!
//! # Storage Layout
//!
//! | Key Pattern | Value | Purpose |
//! |-------------|-------|---------|
//! | `\x00PALS:{canonical}` | JSON-serialized `StoredPredicateAliasSet` | Alias set by canonical name |
//!
//! # Merge Semantics
//!
//! When importing aliases from multiple Trust Packs, aliases are merged:
//! - If Pack A defines `enabled ↔ required`
//! - And Pack B defines `enabled ↔ mandatory`
//! - The result is `enabled ↔ [required, mandatory]`
//!
//! This allows organizations to compose policies without conflict.

use crate::error::Result;
use crate::key_codec;
use crate::traits::KVStore;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tracing::{debug, instrument};

/// A stored predicate alias set.
///
/// This is the serializable form persisted to storage.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct StoredPredicateAliasSet {
    /// The canonical predicate name (used as the index key).
    pub canonical: String,
    /// All aliases that map to this canonical name.
    pub aliases: Vec<String>,
}

impl StoredPredicateAliasSet {
    /// Create a new stored predicate alias set.
    pub fn new(canonical: impl Into<String>, aliases: Vec<impl Into<String>>) -> Self {
        Self { canonical: canonical.into(), aliases: aliases.into_iter().map(Into::into).collect() }
    }

    /// Merge another alias set into this one.
    ///
    /// Adds any aliases from `other` that aren't already present.
    /// Both sets must have the same canonical name.
    pub fn merge(&mut self, other: &StoredPredicateAliasSet) {
        debug_assert_eq!(
            self.canonical, other.canonical,
            "Cannot merge alias sets with different canonical names"
        );

        for alias in &other.aliases {
            if !self.aliases.contains(alias) {
                self.aliases.push(alias.clone());
            }
        }
    }
}

/// Specialized storage trait for predicate alias persistence.
///
/// Enables semantic predicate matching to survive restarts.
///
/// # Example
///
/// ```ignore
/// let predicate_alias_store = GenericPredicateAliasStore::new(kv_store);
///
/// // Store alias set during policy import
/// let alias_set = StoredPredicateAliasSet::new("enabled", vec!["required", "mandatory"]);
/// predicate_alias_store.set_predicate_alias_set(&alias_set).await?;
///
/// // Load on startup
/// let all_aliases = predicate_alias_store.list_all_predicate_aliases().await?;
/// ```
#[async_trait]
pub trait PredicateAliasStore: Send + Sync {
    /// Store or merge a predicate alias set.
    ///
    /// If an alias set with the same canonical name already exists,
    /// the aliases are merged (union of both alias lists).
    ///
    /// # Arguments
    /// * `set` - The alias set to store
    async fn set_predicate_alias_set(&self, set: &StoredPredicateAliasSet) -> Result<()>;

    /// Get a predicate alias set by canonical name.
    ///
    /// Returns the stored alias set if found, or None if not present.
    ///
    /// # Arguments
    /// * `canonical` - The canonical predicate name
    async fn get_predicate_alias_set(
        &self,
        canonical: &str,
    ) -> Result<Option<StoredPredicateAliasSet>>;

    /// List all stored predicate alias sets.
    ///
    /// Used on startup to load persisted aliases into memory.
    async fn list_all_predicate_aliases(&self) -> Result<Vec<StoredPredicateAliasSet>>;
}

/// PredicateAliasStore implementation backed by a generic KVStore.
///
/// Uses JSON serialization for simplicity and debuggability.
pub struct GenericPredicateAliasStore<S> {
    store: S,
}

impl<S: KVStore> GenericPredicateAliasStore<S> {
    /// Create a new PredicateAliasStore backed by the given KVStore.
    pub fn new(store: S) -> Self {
        Self { store }
    }
}

#[async_trait]
impl<S: KVStore + 'static> PredicateAliasStore for GenericPredicateAliasStore<S> {
    #[instrument(skip(self, set), fields(canonical = %set.canonical, alias_count = set.aliases.len()))]
    async fn set_predicate_alias_set(&self, set: &StoredPredicateAliasSet) -> Result<()> {
        let key = key_codec::predicate_alias_key(&set.canonical);

        // Check for existing set to merge with
        let merged_set = if let Some(existing_bytes) = self.store.get(&key).await? {
            let mut existing: StoredPredicateAliasSet = serde_json::from_slice(&existing_bytes)
                .map_err(|e| {
                    crate::error::StorageError::Serialization(format!(
                        "Failed to deserialize predicate alias set: {}",
                        e
                    ))
                })?;

            existing.merge(set);
            debug!(
                canonical = %set.canonical,
                total_aliases = existing.aliases.len(),
                "Merged predicate alias set"
            );
            existing
        } else {
            debug!(canonical = %set.canonical, "Created new predicate alias set");
            set.clone()
        };

        // Serialize to JSON
        let value = serde_json::to_vec(&merged_set).map_err(|e| {
            crate::error::StorageError::Serialization(format!(
                "Failed to serialize predicate alias set: {}",
                e
            ))
        })?;

        self.store.put(&key, &value).await?;
        Ok(())
    }

    #[instrument(skip(self), fields(canonical = %canonical))]
    async fn get_predicate_alias_set(
        &self,
        canonical: &str,
    ) -> Result<Option<StoredPredicateAliasSet>> {
        let key = key_codec::predicate_alias_key(canonical);

        match self.store.get(&key).await? {
            Some(value) => {
                let set: StoredPredicateAliasSet = serde_json::from_slice(&value).map_err(|e| {
                    crate::error::StorageError::Serialization(format!(
                        "Failed to deserialize predicate alias set: {}",
                        e
                    ))
                })?;
                debug!(canonical, alias_count = set.aliases.len(), "Found predicate alias set");
                Ok(Some(set))
            }
            None => {
                debug!(canonical, "No predicate alias set found");
                Ok(None)
            }
        }
    }

    #[instrument(skip(self))]
    async fn list_all_predicate_aliases(&self) -> Result<Vec<StoredPredicateAliasSet>> {
        let prefix = key_codec::predicate_alias_scan_prefix();
        let entries = self.store.scan_prefix(&prefix).await?;

        let mut sets = Vec::with_capacity(entries.len());
        for (_key, value) in entries {
            let set: StoredPredicateAliasSet = serde_json::from_slice(&value).map_err(|e| {
                crate::error::StorageError::Serialization(format!(
                    "Failed to deserialize predicate alias set: {}",
                    e
                ))
            })?;
            sets.push(set);
        }

        debug!(count = sets.len(), "Listed all predicate alias sets");
        Ok(sets)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::HybridStore;
    use std::sync::Arc;

    #[tokio::test]
    async fn test_set_and_get_predicate_alias_set() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        let set = StoredPredicateAliasSet::new("enabled", vec!["required", "mandatory"]);

        // Store alias set
        alias_store.set_predicate_alias_set(&set).await.expect("set");

        // Retrieve alias set
        let retrieved = alias_store.get_predicate_alias_set("enabled").await.expect("get");
        assert!(retrieved.is_some());

        let retrieved_set = retrieved.expect("set");
        assert_eq!(retrieved_set.canonical, "enabled");
        assert!(retrieved_set.aliases.contains(&"required".to_string()));
        assert!(retrieved_set.aliases.contains(&"mandatory".to_string()));
    }

    #[tokio::test]
    async fn test_merge_on_set() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        // First set: enabled ↔ required
        let set1 = StoredPredicateAliasSet::new("enabled", vec!["required"]);
        alias_store.set_predicate_alias_set(&set1).await.expect("set 1");

        // Second set: enabled ↔ mandatory (should merge)
        let set2 = StoredPredicateAliasSet::new("enabled", vec!["mandatory"]);
        alias_store.set_predicate_alias_set(&set2).await.expect("set 2");

        // Should have both aliases
        let retrieved = alias_store.get_predicate_alias_set("enabled").await.expect("get");
        let retrieved_set = retrieved.expect("set");
        assert_eq!(retrieved_set.aliases.len(), 2);
        assert!(retrieved_set.aliases.contains(&"required".to_string()));
        assert!(retrieved_set.aliases.contains(&"mandatory".to_string()));
    }

    #[tokio::test]
    async fn test_merge_deduplicates() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        // Store same alias twice
        let set1 = StoredPredicateAliasSet::new("enabled", vec!["required"]);
        alias_store.set_predicate_alias_set(&set1).await.expect("set 1");

        let set2 = StoredPredicateAliasSet::new("enabled", vec!["required"]); // Same alias
        alias_store.set_predicate_alias_set(&set2).await.expect("set 2");

        // Should only have one "required"
        let retrieved = alias_store.get_predicate_alias_set("enabled").await.expect("get");
        let retrieved_set = retrieved.expect("set");
        assert_eq!(retrieved_set.aliases.len(), 1);
    }

    #[tokio::test]
    async fn test_get_nonexistent() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        let result = alias_store.get_predicate_alias_set("nonexistent").await.expect("get");
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_list_all() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        // Store multiple alias sets
        alias_store
            .set_predicate_alias_set(&StoredPredicateAliasSet::new("enabled", vec!["required"]))
            .await
            .expect("set 1");

        alias_store
            .set_predicate_alias_set(&StoredPredicateAliasSet::new("allowed", vec!["permitted"]))
            .await
            .expect("set 2");

        // List all
        let all_sets = alias_store.list_all_predicate_aliases().await.expect("list");
        assert_eq!(all_sets.len(), 2);

        let canonicals: Vec<_> = all_sets.iter().map(|s| s.canonical.as_str()).collect();
        assert!(canonicals.contains(&"enabled"));
        assert!(canonicals.contains(&"allowed"));
    }

    #[tokio::test]
    async fn test_list_empty() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        let all_sets = alias_store.list_all_predicate_aliases().await.expect("list");
        assert!(all_sets.is_empty());
    }

    #[tokio::test]
    async fn test_different_canonicals_isolated() {
        let store = Arc::new(HybridStore::open_temp().expect("Failed to create store"));
        let alias_store = GenericPredicateAliasStore::new(store);

        // Store two different canonical names
        alias_store
            .set_predicate_alias_set(&StoredPredicateAliasSet::new("enabled", vec!["required"]))
            .await
            .expect("set 1");

        alias_store
            .set_predicate_alias_set(&StoredPredicateAliasSet::new(
                "allowed",
                vec!["permitted", "approved"],
            ))
            .await
            .expect("set 2");

        // Each should be isolated
        let enabled = alias_store.get_predicate_alias_set("enabled").await.expect("get 1");
        assert_eq!(enabled.expect("enabled").aliases.len(), 1);

        let allowed = alias_store.get_predicate_alias_set("allowed").await.expect("get 2");
        assert_eq!(allowed.expect("allowed").aliases.len(), 2);
    }
}
