rustytorch/crates/specialized/rtx-platform/src/tenant.rs

//! Per-tenant isolation and resource quota management module

use crate::error::TenantError;
use crate::{PlatformError, PlatformResult};
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, broadcast};
use uuid::Uuid;

/// Tenant isolation levels
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum IsolationLevel {
    /// No resource sharing, dedicated hardware
    Strict,
    /// Managed sharing with strong isolation guarantees
    Standard,
    /// Best-effort isolation, allows resource sharing
    Shared,
}

/// Resource quota configuration
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ResourceQuota {
    /// Hard limit for this resource
    pub limit: u64,
    /// Currently used amount
    pub used: u64,
    /// Reserved but not yet used
    pub reserved: u64,
    /// Burst limit (temporary overcommit allowed)
    pub burst_limit: Option<u64>,
}

impl ResourceQuota {
    /// Check if allocation is within quota
    pub fn can_allocate(&self, amount: u64) -> bool {
        self.used + self.reserved + amount <= self.limit
    }

    /// Check if burst allocation is possible
    pub fn can_burst_allocate(&self, amount: u64) -> bool {
        if let Some(burst_limit) = self.burst_limit {
            self.used + self.reserved + amount <= burst_limit
        } else {
            false
        }
    }

    /// Allocate resources
    pub fn allocate(&mut self, amount: u64) -> Result<(), TenantError> {
        if !self.can_allocate(amount) {
            return Err(TenantError::QuotaExceeded {
                tenant_id: Uuid::new_v4(), // This will be filled by caller
                resource: "generic".to_string(),
            });
        }
        self.used += amount;
        Ok(())
    }

    /// Reserve resources
    pub fn reserve(&mut self, amount: u64) -> Result<(), TenantError> {
        if self.used + self.reserved + amount > self.limit {
            return Err(TenantError::QuotaExceeded {
                tenant_id: Uuid::new_v4(),
                resource: "generic".to_string(),
            });
        }
        self.reserved += amount;
        Ok(())
    }

    /// Activate reservation (move from reserved to used)
    pub fn activate_reservation(&mut self, amount: u64) {
        let actual_amount = self.reserved.min(amount);
        self.reserved -= actual_amount;
        self.used += actual_amount;
    }

    /// Release resources
    pub fn release(&mut self, amount: u64) {
        self.used = self.used.saturating_sub(amount);
    }
}

/// Tenant configuration
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TenantConfig {
    pub id: Uuid,
    pub name: String,
    pub created_at: DateTime<Utc>,
    pub isolation_level: IsolationLevel,
    pub resource_quotas: HashMap<String, ResourceQuota>,
    pub allowed_regions: Vec<String>,
    pub priority: u32,
}

/// Tenant status information
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum TenantStatus {
    Active,
    Suspended,
    Terminated,
}

/// Resource allocation result
#[derive(Debug, Clone)]
pub struct ResourceAllocation {
    pub allocation_id: Uuid,
    pub tenant_id: Uuid,
    pub region: String,
    pub allocated_resources: HashMap<String, u64>,
    pub is_burst_allocation: bool,
    pub burst_expires_at: Option<DateTime<Utc>>,
    pub allocated_at: DateTime<Utc>,
}

/// API key for tenant authentication
#[derive(Debug, Clone)]
pub struct ApiKey {
    pub key: String,
    pub tenant_id: Uuid,
    pub permissions: Vec<String>,
    pub expires_at: Option<DateTime<Utc>>,
    pub created_at: DateTime<Utc>,
}

/// Authentication context
#[derive(Debug, Clone)]
pub struct AuthContext {
    pub tenant_id: Uuid,
    pub permissions: Vec<String>,
    pub authenticated_at: DateTime<Utc>,
}

/// Resource reservation
#[derive(Debug, Clone)]
pub struct ResourceReservation {
    pub reservation_id: Option<Uuid>,
    pub tenant_id: Uuid,
    pub region: String,
    pub resources: HashMap<String, u64>,
    pub scheduled_start: DateTime<Utc>,
    pub duration: chrono::Duration,
    pub created_at: DateTime<Utc>,
}

/// Isolation check result
#[derive(Debug, Clone)]
pub struct IsolationCheck {
    pub tenant_id: Uuid,
    pub violations: Vec<String>,
    pub has_dedicated_resources: bool,
    pub checked_at: DateTime<Utc>,
}

/// Tenant manager for isolation and resource management
#[derive(Debug)]
pub struct TenantManager {
    tenants: Arc<DashMap<Uuid, TenantConfig>>,
    tenant_status: Arc<DashMap<Uuid, TenantStatus>>,
    allocations: Arc<DashMap<Uuid, ResourceAllocation>>,
    api_keys: Arc<DashMap<String, ApiKey>>,
    reservations: Arc<DashMap<Uuid, ResourceReservation>>,
    global_quotas: Arc<RwLock<HashMap<Uuid, HashMap<String, u64>>>>, // tenant -> resource -> global_used
    shutdown_tx: Option<broadcast::Sender<()>>,
    redis_client: Option<redis::Client>,
}

impl TenantManager {
    /// Create new TenantManager
    pub async fn new(config: &crate::PlatformConfig) -> PlatformResult<Self> {
        let tenants = Arc::new(DashMap::new());
        let tenant_status = Arc::new(DashMap::new());
        let allocations = Arc::new(DashMap::new());
        let api_keys = Arc::new(DashMap::new());
        let reservations = Arc::new(DashMap::new());
        let global_quotas = Arc::new(RwLock::new(HashMap::new()));

        // Initialize Redis client for distributed state
        let redis_client = if !config.redis_urls.is_empty() {
            Some(redis::Client::open(config.redis_urls[0].as_str()).map_err(PlatformError::Redis)?)
        } else {
            None
        };

        Ok(Self {
            tenants,
            tenant_status,
            allocations,
            api_keys,
            reservations,
            global_quotas,
            shutdown_tx: None,
            redis_client,
        })
    }

    /// Start tenant manager services
    pub async fn start(&mut self) -> PlatformResult<()> {
        let (shutdown_tx, mut shutdown_rx) = broadcast::channel(1);
        self.shutdown_tx = Some(shutdown_tx);

        // Start background tasks
        let global_quotas = self.global_quotas.clone();
        let allocations = self.allocations.clone();

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60));

            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Clean up expired burst allocations
                        Self::cleanup_expired_allocations(&allocations).await;

                        // Update global quota tracking
                        Self::update_global_quotas(&global_quotas, &allocations).await;
                    }
                    _ = shutdown_rx.recv() => {
                        break;
                    }
                }
            }
        });

        tracing::info!("TenantManager started");
        Ok(())
    }

    /// Shutdown tenant manager
    pub async fn shutdown(&mut self) -> PlatformResult<()> {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
        tracing::info!("TenantManager shutdown");
        Ok(())
    }

    /// Create a new tenant
    pub async fn create_tenant(&mut self, config: TenantConfig) -> PlatformResult<()> {
        let tenant_id = config.id;

        // Initialize global quota tracking
        let mut global_quotas = self.global_quotas.write().await;
        let mut global_tenant_quotas = HashMap::new();
        for resource_type in config.resource_quotas.keys() {
            global_tenant_quotas.insert(resource_type.clone(), 0u64);
        }
        global_quotas.insert(tenant_id, global_tenant_quotas);
        drop(global_quotas);

        self.tenants.insert(tenant_id, config);
        self.tenant_status.insert(tenant_id, TenantStatus::Active);

        tracing::info!("Created tenant: {}", tenant_id);
        Ok(())
    }

    /// Get tenant configuration
    pub async fn get_tenant(&self, tenant_id: Uuid) -> PlatformResult<TenantConfig> {
        let mut tenant = self
            .tenants
            .get(&tenant_id)
            .map(|t| t.clone())
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        // Update resource usage from allocations
        for allocation_entry in self.allocations.iter() {
            let allocation = allocation_entry.value();
            if allocation.tenant_id == tenant_id {
                for (resource_type, amount) in &allocation.allocated_resources {
                    if let Some(quota) = tenant.resource_quotas.get_mut(resource_type) {
                        // Note: This is a simplified approach. In production, you'd want
                        // more sophisticated tracking to avoid double-counting
                        quota.used = quota.used.max(*amount);
                    }
                }
            }
        }

        Ok(tenant)
    }

    /// Allocate resources for a tenant
    pub async fn allocate_resources(
        &self,
        tenant_id: Uuid,
        region: &str,
        resources: HashMap<String, u64>,
    ) -> PlatformResult<ResourceAllocation> {
        let mut tenant = self
            .tenants
            .get_mut(&tenant_id)
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        // Check if tenant is allowed in this region
        if !tenant.allowed_regions.contains(&region.to_string()) {
            return Err(PlatformError::Tenant(TenantError::IsolationViolation {
                tenant_id,
            }));
        }

        // Check quotas for each resource
        for (resource_type, requested_amount) in &resources {
            if let Some(quota) = tenant.resource_quotas.get(resource_type)
                && !quota.can_allocate(*requested_amount)
            {
                return Err(PlatformError::Tenant(TenantError::QuotaExceeded {
                    tenant_id,
                    resource: resource_type.clone(),
                }));
            }
        }

        // Update quotas and global tracking
        let mut global_quotas = self.global_quotas.write().await;
        let global_tenant_quotas = global_quotas.entry(tenant_id).or_insert_with(HashMap::new);

        for (resource_type, requested_amount) in &resources {
            // Check global quota across all regions
            let global_used = global_tenant_quotas
                .get(resource_type)
                .copied()
                .unwrap_or(0);
            let quota_limit = tenant
                .resource_quotas
                .get(resource_type)
                .map_or(0, |q| q.limit);

            if global_used + requested_amount > quota_limit {
                return Err(PlatformError::Tenant(TenantError::QuotaExceeded {
                    tenant_id,
                    resource: resource_type.clone(),
                }));
            }

            // Update local quota
            if let Some(quota) = tenant.resource_quotas.get_mut(resource_type) {
                quota.allocate(*requested_amount).map_err(|mut e| {
                    if let TenantError::QuotaExceeded {
                        tenant_id: ref mut tid,
                        ..
                    } = e
                    {
                        *tid = tenant_id;
                    }
                    PlatformError::Tenant(e)
                })?;
            }

            // Update global tracking
            global_tenant_quotas.insert(resource_type.clone(), global_used + requested_amount);
        }

        let allocation = ResourceAllocation {
            allocation_id: Uuid::new_v4(),
            tenant_id,
            region: region.to_string(),
            allocated_resources: resources,
            is_burst_allocation: false,
            burst_expires_at: None,
            allocated_at: Utc::now(),
        };

        self.allocations
            .insert(allocation.allocation_id, allocation.clone());
        Ok(allocation)
    }

    /// Allocate resources with burst capability
    pub async fn allocate_resources_with_burst(
        &self,
        tenant_id: Uuid,
        region: &str,
        resources: HashMap<String, u64>,
        burst_duration: Duration,
    ) -> PlatformResult<ResourceAllocation> {
        let mut tenant = self
            .tenants
            .get_mut(&tenant_id)
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        // Try normal allocation first
        if let Ok(allocation) = self
            .allocate_resources(tenant_id, region, resources.clone())
            .await
        {
            return Ok(allocation);
        }

        // Check if burst allocation is possible
        for (resource_type, requested_amount) in &resources {
            if let Some(quota) = tenant.resource_quotas.get(resource_type)
                && !quota.can_burst_allocate(*requested_amount)
            {
                return Err(PlatformError::Tenant(TenantError::QuotaExceeded {
                    tenant_id,
                    resource: resource_type.clone(),
                }));
            }
        }

        // Perform burst allocation
        for (resource_type, requested_amount) in &resources {
            if let Some(quota) = tenant.resource_quotas.get_mut(resource_type) {
                quota.used += requested_amount;
            }
        }

        let burst_expires_at = Utc::now() + chrono::Duration::from_std(burst_duration).unwrap();

        let allocation = ResourceAllocation {
            allocation_id: Uuid::new_v4(),
            tenant_id,
            region: region.to_string(),
            allocated_resources: resources,
            is_burst_allocation: true,
            burst_expires_at: Some(burst_expires_at),
            allocated_at: Utc::now(),
        };

        self.allocations
            .insert(allocation.allocation_id, allocation.clone());
        Ok(allocation)
    }

    /// Clean up expired burst allocations
    pub async fn cleanup_expired_burst_allocations(&self) -> PlatformResult<()> {
        Self::cleanup_expired_allocations(&self.allocations).await;
        Ok(())
    }

    async fn cleanup_expired_allocations(allocations: &Arc<DashMap<Uuid, ResourceAllocation>>) {
        let now = Utc::now();
        let mut to_remove = Vec::new();

        for allocation_entry in allocations.iter() {
            let allocation = allocation_entry.value();
            if allocation.is_burst_allocation
                && let Some(expires_at) = allocation.burst_expires_at
                && now > expires_at
            {
                to_remove.push(allocation.allocation_id);
            }
        }

        for allocation_id in to_remove {
            allocations.remove(&allocation_id);
        }
    }

    async fn update_global_quotas(
        global_quotas: &Arc<RwLock<HashMap<Uuid, HashMap<String, u64>>>>,
        allocations: &Arc<DashMap<Uuid, ResourceAllocation>>,
    ) {
        let mut quotas = global_quotas.write().await;
        quotas.clear();

        // Recalculate global usage from active allocations
        for allocation_entry in allocations.iter() {
            let allocation = allocation_entry.value();
            let tenant_quotas = quotas
                .entry(allocation.tenant_id)
                .or_insert_with(HashMap::new);

            for (resource_type, amount) in &allocation.allocated_resources {
                let current = tenant_quotas.get(resource_type).copied().unwrap_or(0);
                tenant_quotas.insert(resource_type.clone(), current + amount);
            }
        }
    }

    /// Check isolation violations
    pub async fn check_isolation_violations(
        &self,
        tenant_id: Uuid,
    ) -> PlatformResult<IsolationCheck> {
        let tenant = self
            .tenants
            .get(&tenant_id)
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        let violations = Vec::new(); // No violations in current simple implementation
        let has_dedicated_resources = tenant.isolation_level == IsolationLevel::Strict;

        Ok(IsolationCheck {
            tenant_id,
            violations,
            has_dedicated_resources,
            checked_at: Utc::now(),
        })
    }

    /// Attempt cross-tenant access (should fail for strict isolation)
    pub async fn attempt_cross_tenant_access(
        &self,
        requesting_tenant: Uuid,
        target_tenant: Uuid,
        _operation: &str,
    ) -> PlatformResult<()> {
        if requesting_tenant == target_tenant {
            return Ok(());
        }

        let target_tenant_config = self.tenants.get(&target_tenant).ok_or_else(|| {
            PlatformError::Tenant(TenantError::NotFound {
                tenant_id: target_tenant,
            })
        })?;

        // Strict isolation prevents any cross-tenant access
        if target_tenant_config.isolation_level == IsolationLevel::Strict {
            return Err(PlatformError::Tenant(TenantError::IsolationViolation {
                tenant_id: requesting_tenant,
            }));
        }

        Ok(())
    }

    /// Generate API key for tenant
    pub async fn generate_tenant_api_key(
        &self,
        tenant_id: Uuid,
        permissions: Vec<String>,
        expires_at: Option<DateTime<Utc>>,
    ) -> PlatformResult<ApiKey> {
        // Verify tenant exists
        self.tenants
            .get(&tenant_id)
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        // Generate secure API key
        let key_bytes: [u8; 32] = rand::random();
        let key = format!("rtx_{}", hex::encode(key_bytes));

        let api_key = ApiKey {
            key: key.clone(),
            tenant_id,
            permissions,
            expires_at,
            created_at: Utc::now(),
        };

        self.api_keys.insert(key, api_key.clone());
        Ok(api_key)
    }

    /// Authenticate request with API key
    pub async fn authenticate_request(
        &self,
        api_key: &str,
        required_permission: &str,
    ) -> PlatformResult<AuthContext> {
        let key_info = self.api_keys.get(api_key).ok_or_else(|| {
            PlatformError::Tenant(TenantError::AuthFailed {
                tenant_id: Uuid::nil(),
            })
        })?;

        // Check if key is expired
        if let Some(expires_at) = key_info.expires_at
            && Utc::now() > expires_at
        {
            return Err(PlatformError::Tenant(TenantError::AuthFailed {
                tenant_id: key_info.tenant_id,
            }));
        }

        // Check permissions
        if !key_info
            .permissions
            .contains(&required_permission.to_string())
        {
            return Err(PlatformError::Tenant(TenantError::AuthFailed {
                tenant_id: key_info.tenant_id,
            }));
        }

        Ok(AuthContext {
            tenant_id: key_info.tenant_id,
            permissions: key_info.permissions.clone(),
            authenticated_at: Utc::now(),
        })
    }

    /// Create resource reservation
    pub async fn create_resource_reservation(
        &self,
        tenant_id: Uuid,
        region: &str,
        resources: HashMap<String, u64>,
        scheduled_start: DateTime<Utc>,
        duration: chrono::Duration,
    ) -> PlatformResult<ResourceReservation> {
        let mut tenant = self
            .tenants
            .get_mut(&tenant_id)
            .ok_or_else(|| PlatformError::Tenant(TenantError::NotFound { tenant_id }))?;

        // Check if tenant can reserve these resources
        for (resource_type, requested_amount) in &resources {
            if let Some(quota) = tenant.resource_quotas.get_mut(resource_type) {
                quota.reserve(*requested_amount).map_err(|mut e| {
                    if let TenantError::QuotaExceeded {
                        tenant_id: ref mut tid,
                        resource: ref mut res,
                    } = e
                    {
                        *tid = tenant_id;
                        *res = resource_type.clone();
                    }
                    PlatformError::Tenant(e)
                })?;
            }
        }

        let reservation_id = Uuid::new_v4();
        let reservation = ResourceReservation {
            reservation_id: Some(reservation_id),
            tenant_id,
            region: region.to_string(),
            resources,
            scheduled_start,
            duration,
            created_at: Utc::now(),
        };

        self.reservations
            .insert(reservation_id, reservation.clone());
        Ok(reservation)
    }

    /// Activate a reservation (convert reserved to used)
    pub async fn activate_reservation(&self, reservation_id: Uuid) -> PlatformResult<()> {
        let reservation =
            self.reservations
                .get(&reservation_id)
                .ok_or_else(|| PlatformError::Internal {
                    message: "Reservation not found".to_string(),
                })?;

        let mut tenant = self
            .tenants
            .get_mut(&reservation.tenant_id)
            .ok_or_else(|| {
                PlatformError::Tenant(TenantError::NotFound {
                    tenant_id: reservation.tenant_id,
                })
            })?;

        // Activate reservation for each resource
        for (resource_type, amount) in &reservation.resources {
            if let Some(quota) = tenant.resource_quotas.get_mut(resource_type) {
                quota.activate_reservation(*amount);
            }
        }

        Ok(())
    }
}

// Required for hex encoding in API key generation
mod hex {
    pub fn encode(bytes: [u8; 32]) -> String {
        bytes.iter().map(|b| format!("{b:02x}")).collect()
    }
}

// Simple random number generation for API keys
mod rand {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    use std::time::{SystemTime, UNIX_EPOCH};

    pub fn random() -> [u8; 32] {
        let mut hasher = DefaultHasher::new();
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
            .hash(&mut hasher);

        let hash = hasher.finish();
        let mut bytes = [0u8; 32];
        for (i, byte) in bytes.iter_mut().enumerate() {
            *byte = ((hash >> (i * 8)) & 0xFF) as u8;
        }
        bytes
    }
}