enact-core 0.0.2

//! Tenant Policy - Multi-tenant isolation and limits

use super::{PolicyContext, PolicyDecision, PolicyEvaluator};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// Tenant limits
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TenantLimits {
    /// Maximum concurrent executions
    pub max_concurrent_executions: Option<usize>,
    /// Maximum executions per day
    pub max_executions_per_day: Option<usize>,
    /// Maximum LLM tokens per day
    pub max_tokens_per_day: Option<usize>,
    /// Maximum storage in bytes
    pub max_storage_bytes: Option<usize>,
}

impl Default for TenantLimits {
    fn default() -> Self {
        Self {
            max_concurrent_executions: Some(10),
            max_executions_per_day: Some(1000),
            max_tokens_per_day: Some(1_000_000),
            max_storage_bytes: Some(1024 * 1024 * 1024), // 1GB
        }
    }
}

impl TenantLimits {
    /// Create unlimited tenant limits (for enterprise/unlimited tier)
    pub fn unlimited() -> Self {
        Self {
            max_concurrent_executions: None,
            max_executions_per_day: None,
            max_tokens_per_day: None,
            max_storage_bytes: None,
        }
    }

    /// Create free tier limits
    pub fn free_tier() -> Self {
        Self {
            max_concurrent_executions: Some(2),
            max_executions_per_day: Some(100),
            max_tokens_per_day: Some(50_000),
            max_storage_bytes: Some(100 * 1024 * 1024), // 100MB
        }
    }

    /// Create pro tier limits
    pub fn pro_tier() -> Self {
        Self {
            max_concurrent_executions: Some(20),
            max_executions_per_day: Some(10_000),
            max_tokens_per_day: Some(10_000_000),
            max_storage_bytes: Some(10 * 1024 * 1024 * 1024), // 10GB
        }
    }
}

/// Feature flags for tenant
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct FeatureFlags {
    /// Enabled features
    pub enabled: HashSet<String>,
    /// Disabled features (takes precedence)
    pub disabled: HashSet<String>,
}

impl FeatureFlags {
    /// Create new feature flags
    pub fn new() -> Self {
        Self::default()
    }

    /// Enable a feature
    pub fn enable(mut self, feature: impl Into<String>) -> Self {
        self.enabled.insert(feature.into());
        self
    }

    /// Disable a feature
    pub fn disable(mut self, feature: impl Into<String>) -> Self {
        self.disabled.insert(feature.into());
        self
    }

    /// Check if a feature is enabled
    pub fn is_enabled(&self, feature: &str) -> bool {
        !self.disabled.contains(feature) && self.enabled.contains(feature)
    }

    /// Common feature flags
    pub fn with_defaults() -> Self {
        Self::new()
            .enable("basic_execution")
            .enable("tool_invocation")
            .enable("streaming")
    }

    /// All features enabled
    pub fn all_enabled() -> Self {
        Self::new()
            .enable("basic_execution")
            .enable("tool_invocation")
            .enable("streaming")
            .enable("parallel_execution")
            .enable("nested_execution")
            .enable("custom_tools")
            .enable("mcp_integration")
            .enable("advanced_memory")
    }
}

/// Tenant policy
#[derive(Debug, Clone)]
pub struct TenantPolicy {
    /// Tenant limits
    pub limits: TenantLimits,
    /// Feature flags
    pub features: FeatureFlags,
    /// Allowed models
    pub allowed_models: HashSet<String>,
    /// Allowed tools
    pub allowed_tools: Option<HashSet<String>>,
    /// Isolation mode
    pub isolation: TenantIsolation,
}

/// Tenant isolation mode
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum TenantIsolation {
    /// Shared resources (default)
    #[default]
    Shared,
    /// Dedicated resources
    Dedicated,
    /// Strict isolation (no shared state)
    Strict,
}

impl Default for TenantPolicy {
    fn default() -> Self {
        Self {
            limits: TenantLimits::default(),
            features: FeatureFlags::with_defaults(),
            allowed_models: HashSet::new(),
            allowed_tools: None, // None means all tools allowed
            isolation: TenantIsolation::default(),
        }
    }
}

impl TenantPolicy {
    /// Create a new tenant policy
    pub fn new() -> Self {
        Self::default()
    }

    /// Set limits
    pub fn with_limits(mut self, limits: TenantLimits) -> Self {
        self.limits = limits;
        self
    }

    /// Set features
    pub fn with_features(mut self, features: FeatureFlags) -> Self {
        self.features = features;
        self
    }

    /// Allow a specific model
    pub fn allow_model(mut self, model: impl Into<String>) -> Self {
        self.allowed_models.insert(model.into());
        self
    }

    /// Allow a specific tool
    pub fn allow_tool(mut self, tool: impl Into<String>) -> Self {
        self.allowed_tools
            .get_or_insert_with(HashSet::new)
            .insert(tool.into());
        self
    }

    /// Set isolation mode
    pub fn with_isolation(mut self, isolation: TenantIsolation) -> Self {
        self.isolation = isolation;
        self
    }

    /// Check if a model is allowed
    pub fn is_model_allowed(&self, model: &str) -> bool {
        self.allowed_models.is_empty() || self.allowed_models.contains(model)
    }

    /// Check if a tool is allowed
    pub fn is_tool_allowed(&self, tool: &str) -> bool {
        self.allowed_tools
            .as_ref()
            .map(|tools| tools.contains(tool))
            .unwrap_or(true)
    }
}

impl PolicyEvaluator for TenantPolicy {
    fn evaluate(&self, context: &PolicyContext) -> PolicyDecision {
        match &context.action {
            super::PolicyAction::LlmCall { model } => {
                if !self.is_model_allowed(model) {
                    return PolicyDecision::Deny {
                        reason: format!("Model '{}' is not allowed for this tenant", model),
                    };
                }
                PolicyDecision::Allow
            }
            super::PolicyAction::InvokeTool { tool_name } => {
                if !self.is_tool_allowed(tool_name) {
                    return PolicyDecision::Deny {
                        reason: format!("Tool '{}' is not allowed for this tenant", tool_name),
                    };
                }
                PolicyDecision::Allow
            }
            _ => PolicyDecision::Allow,
        }
    }
}

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

    // ============ TenantLimits Tests ============

    #[test]
    fn test_tenant_limits_default() {
        let limits = TenantLimits::default();
        assert_eq!(limits.max_concurrent_executions, Some(10));
        assert_eq!(limits.max_executions_per_day, Some(1000));
        assert_eq!(limits.max_tokens_per_day, Some(1_000_000));
        assert_eq!(limits.max_storage_bytes, Some(1024 * 1024 * 1024));
    }

    #[test]
    fn test_tenant_limits_unlimited() {
        let limits = TenantLimits::unlimited();
        assert!(limits.max_concurrent_executions.is_none());
        assert!(limits.max_executions_per_day.is_none());
        assert!(limits.max_tokens_per_day.is_none());
        assert!(limits.max_storage_bytes.is_none());
    }

    #[test]
    fn test_tenant_limits_free_tier() {
        let limits = TenantLimits::free_tier();
        assert_eq!(limits.max_concurrent_executions, Some(2));
        assert_eq!(limits.max_executions_per_day, Some(100));
        assert_eq!(limits.max_tokens_per_day, Some(50_000));
    }

    #[test]
    fn test_tenant_limits_pro_tier() {
        let limits = TenantLimits::pro_tier();
        assert_eq!(limits.max_concurrent_executions, Some(20));
        assert_eq!(limits.max_executions_per_day, Some(10_000));
        assert_eq!(limits.max_tokens_per_day, Some(10_000_000));
    }

    // ============ FeatureFlags Tests ============

    #[test]
    fn test_feature_flags_default() {
        let flags = FeatureFlags::default();
        assert!(flags.enabled.is_empty());
        assert!(flags.disabled.is_empty());
    }

    #[test]
    fn test_feature_flags_enable() {
        let flags = FeatureFlags::new().enable("feature_a").enable("feature_b");

        assert!(flags.is_enabled("feature_a"));
        assert!(flags.is_enabled("feature_b"));
        assert!(!flags.is_enabled("feature_c"));
    }

    #[test]
    fn test_feature_flags_disable_overrides_enable() {
        let flags = FeatureFlags::new().enable("feature_x").disable("feature_x");

        // Disabled takes precedence
        assert!(!flags.is_enabled("feature_x"));
    }

    #[test]
    fn test_feature_flags_with_defaults() {
        let flags = FeatureFlags::with_defaults();
        assert!(flags.is_enabled("basic_execution"));
        assert!(flags.is_enabled("tool_invocation"));
        assert!(flags.is_enabled("streaming"));
    }

    #[test]
    fn test_feature_flags_all_enabled() {
        let flags = FeatureFlags::all_enabled();
        assert!(flags.is_enabled("basic_execution"));
        assert!(flags.is_enabled("parallel_execution"));
        assert!(flags.is_enabled("mcp_integration"));
        assert!(flags.is_enabled("advanced_memory"));
    }

    // ============ TenantIsolation Tests ============

    #[test]
    fn test_tenant_isolation_default() {
        assert_eq!(TenantIsolation::default(), TenantIsolation::Shared);
    }

    // ============ TenantPolicy Tests ============

    #[test]
    fn test_tenant_policy_default() {
        let policy = TenantPolicy::default();
        assert!(policy.allowed_models.is_empty());
        assert!(policy.allowed_tools.is_none());
        assert_eq!(policy.isolation, TenantIsolation::Shared);
    }

    #[test]
    fn test_tenant_policy_with_limits() {
        let policy = TenantPolicy::new().with_limits(TenantLimits::free_tier());
        assert_eq!(policy.limits.max_concurrent_executions, Some(2));
    }

    #[test]
    fn test_tenant_policy_with_features() {
        let policy = TenantPolicy::new().with_features(FeatureFlags::all_enabled());
        assert!(policy.features.is_enabled("parallel_execution"));
    }

    #[test]
    fn test_tenant_policy_allow_model() {
        let policy = TenantPolicy::new()
            .allow_model("gpt-4")
            .allow_model("claude-3");

        assert!(policy.is_model_allowed("gpt-4"));
        assert!(policy.is_model_allowed("claude-3"));
        assert!(!policy.is_model_allowed("unknown-model"));
    }

    #[test]
    fn test_tenant_policy_allow_model_empty_allows_all() {
        let policy = TenantPolicy::new();
        // Empty allowed_models means all models are allowed
        assert!(policy.is_model_allowed("any-model"));
    }

    #[test]
    fn test_tenant_policy_allow_tool() {
        let policy = TenantPolicy::new()
            .allow_tool("web_search")
            .allow_tool("calculator");

        assert!(policy.is_tool_allowed("web_search"));
        assert!(policy.is_tool_allowed("calculator"));
        assert!(!policy.is_tool_allowed("file_system"));
    }

    #[test]
    fn test_tenant_policy_no_tool_restriction_allows_all() {
        let policy = TenantPolicy::new();
        // None for allowed_tools means all tools are allowed
        assert!(policy.is_tool_allowed("any-tool"));
    }

    #[test]
    fn test_tenant_policy_with_isolation() {
        let policy = TenantPolicy::new().with_isolation(TenantIsolation::Strict);
        assert_eq!(policy.isolation, TenantIsolation::Strict);
    }

    // ============ PolicyEvaluator Tests ============

    #[test]
    fn test_tenant_policy_evaluate_llm_allowed() {
        let policy = TenantPolicy::new(); // Empty allowed_models = all allowed
        let context = PolicyContext {
            tenant_id: Some("tenant-1".to_string()),
            user_id: None,
            action: super::super::PolicyAction::LlmCall {
                model: "gpt-4".to_string(),
            },
            metadata: HashMap::new(),
        };

        let decision = policy.evaluate(&context);
        assert!(decision.is_allowed());
    }

    #[test]
    fn test_tenant_policy_evaluate_llm_denied() {
        let policy = TenantPolicy::new().allow_model("claude-3"); // Only claude allowed
        let context = PolicyContext {
            tenant_id: Some("tenant-1".to_string()),
            user_id: None,
            action: super::super::PolicyAction::LlmCall {
                model: "gpt-4".to_string(),
            },
            metadata: HashMap::new(),
        };

        let decision = policy.evaluate(&context);
        assert!(decision.is_denied());
    }

    #[test]
    fn test_tenant_policy_evaluate_tool_allowed() {
        let policy = TenantPolicy::new(); // None = all tools allowed
        let context = PolicyContext {
            tenant_id: Some("tenant-1".to_string()),
            user_id: None,
            action: super::super::PolicyAction::InvokeTool {
                tool_name: "any_tool".to_string(),
            },
            metadata: HashMap::new(),
        };

        let decision = policy.evaluate(&context);
        assert!(decision.is_allowed());
    }

    #[test]
    fn test_tenant_policy_evaluate_tool_denied() {
        let policy = TenantPolicy::new().allow_tool("calculator"); // Only calculator allowed
        let context = PolicyContext {
            tenant_id: Some("tenant-1".to_string()),
            user_id: None,
            action: super::super::PolicyAction::InvokeTool {
                tool_name: "web_search".to_string(),
            },
            metadata: HashMap::new(),
        };

        let decision = policy.evaluate(&context);
        assert!(decision.is_denied());
    }
}