//! Federation - Main entry point for heterogeneous storage federation
//!
//! Provides unified access to multiple heterogeneous stores (graph, document, vector)
//! through a single interface.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};

use super::error::{FederationError, Result};
use super::traits::Store;
use super::types::{
    StoreType, Record, WriteOp, BatchWriteResult, UniversalQuery, CrossDbQuery, QueryResult, QueryStats,
};
use super::adapters::{GraphStoreAdapter, DocumentStoreAdapter, ConcurrentDocStoreAdapter, VectorStoreAdapter};
use super::cross_db_index::CrossDbIndex;
use super::query_planner::QueryPlanner;
use super::query_executor::{QueryExecutor, ExecutorConfig};
use super::transaction::TransactionCoordinator;
use super::resilience::ResilienceConfig;
use super::rate_limiter::RateLimitConfig;
use super::reaper::ReaperConfig;
use super::permission::{PermissionChecker, Permission, AccessContext};
use super::record_permission::{RecordPermissionChecker, ProtectionCheckResult};
use super::types::{Visibility, ProtectionLevel, SYSTEM_OWNER_ID};

/// Federation configuration
#[derive(Debug, Clone)]
pub struct FederationConfig {
    /// Enable cross-database indexing
    pub enable_cross_db_index: bool,
    /// Enable query planning/optimization
    pub enable_query_planner: bool,
    /// Index storage path
    pub index_path: Option<String>,
    /// Default query timeout in milliseconds
    pub query_timeout_ms: u64,
    /// Enable transaction logging
    pub enable_tx_logging: bool,
    /// Resilience settings (circuit breaker, retry, timeout)
    pub resilience: ResilienceConfig,
    /// Rate limiting settings
    pub rate_limit: RateLimitConfig,
    /// Transaction reaper settings
    pub reaper: ReaperConfig,
    /// Enable permission checking
    pub enable_permissions: bool,
}

impl Default for FederationConfig {
    fn default() -> Self {
        Self {
            enable_cross_db_index: true,
            enable_query_planner: true,
            index_path: None,
            query_timeout_ms: 30000,
            enable_tx_logging: true,
            resilience: ResilienceConfig::default(),
            rate_limit: RateLimitConfig::default(),
            reaper: ReaperConfig::default(),
            enable_permissions: false,
        }
    }
}

/// Store entry in the federation
struct StoreEntry {
    store: Arc<dyn Store>,
    store_type: StoreType,
}

/// Federation - Unified interface for heterogeneous storage
pub struct Federation {
    /// Configuration
    config: FederationConfig,
    /// Registered stores by alias
    stores: RwLock<HashMap<String, StoreEntry>>,
    /// Graph stores (for type-safe access)
    graph_stores: RwLock<HashMap<String, Arc<GraphStoreAdapter>>>,
    /// Document stores (for type-safe access)
    doc_stores: RwLock<HashMap<String, Arc<DocumentStoreAdapter>>>,
    /// Concurrent document stores (for type-safe access)
    concurrent_doc_stores: RwLock<HashMap<String, Arc<ConcurrentDocStoreAdapter>>>,
    /// Vector stores (for type-safe access)
    vector_stores: RwLock<HashMap<String, Arc<VectorStoreAdapter>>>,
    /// Cross-database index
    index: Arc<CrossDbIndex>,
    /// Query planner
    planner: Arc<QueryPlanner>,
    /// Transaction coordinator
    tx_coordinator: Arc<TransactionCoordinator>,
    /// RBAC permission checker (Layer 1)
    permission_checker: PermissionChecker,
}

impl Federation {
    /// Create a new Federation instance
    pub fn new(config: FederationConfig) -> Self {
        let index = Arc::new(CrossDbIndex::new());
        let planner = Arc::new(QueryPlanner::new());
        let tx_coordinator = Arc::new(TransactionCoordinator::new());
        let permission_checker = if config.enable_permissions {
            PermissionChecker::new()
        } else {
            PermissionChecker::disabled()
        };
        
        Self {
            config,
            stores: RwLock::new(HashMap::new()),
            graph_stores: RwLock::new(HashMap::new()),
            doc_stores: RwLock::new(HashMap::new()),
            concurrent_doc_stores: RwLock::new(HashMap::new()),
            vector_stores: RwLock::new(HashMap::new()),
            index,
            planner,
            tx_coordinator,
            permission_checker,
        }
    }
    
    /// Create with default configuration
    pub fn with_defaults() -> Self {
        Self::new(FederationConfig::default())
    }
    
    // ===== Store Registration =====
    
    /// Register a graph store
    pub fn register_graph(&self, alias: &str, store: Arc<GraphStoreAdapter>) -> Result<()> {
        let mut stores = self.stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        if stores.contains_key(alias) {
            return Err(FederationError::ConfigError(
                format!("Store with alias '{}' already registered", alias)
            ));
        }
        
        stores.insert(alias.to_string(), StoreEntry {
            store: store.clone(),
            store_type: StoreType::Graph,
        });
        
        let mut graph_stores = self.graph_stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        graph_stores.insert(alias.to_string(), store.clone());
        
        // Also register with transaction coordinator for 2PC support
        self.tx_coordinator.register_store(alias, store)?;
        
        Ok(())
    }
    
    /// Register a document store
    pub fn register_document(&self, alias: &str, store: Arc<DocumentStoreAdapter>) -> Result<()> {
        let mut stores = self.stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        if stores.contains_key(alias) {
            return Err(FederationError::ConfigError(
                format!("Store with alias '{}' already registered", alias)
            ));
        }
        
        stores.insert(alias.to_string(), StoreEntry {
            store: store.clone(),
            store_type: StoreType::Document,
        });
        
        let mut doc_stores = self.doc_stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        doc_stores.insert(alias.to_string(), store.clone());
        
        // Also register with transaction coordinator for 2PC support
        self.tx_coordinator.register_store(alias, store)?;
        
        Ok(())
    }
    
    /// Register a concurrent document store (high-concurrency, WAL-based)
    pub fn register_concurrent_document(&self, alias: &str, store: Arc<ConcurrentDocStoreAdapter>) -> Result<()> {
        let mut stores = self.stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        if stores.contains_key(alias) {
            return Err(FederationError::ConfigError(
                format!("Store with alias '{}' already registered", alias)
            ));
        }
        
        stores.insert(alias.to_string(), StoreEntry {
            store: store.clone(),
            store_type: StoreType::Document,
        });
        
        let mut concurrent_doc_stores = self.concurrent_doc_stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        concurrent_doc_stores.insert(alias.to_string(), store.clone());
        
        // Also register with transaction coordinator for 2PC support
        self.tx_coordinator.register_store(alias, store)?;
        
        Ok(())
    }
    
    /// Register a vector store (high-concurrency, sharded HNSW)
    pub fn register_vector(&self, alias: &str, store: Arc<VectorStoreAdapter>) -> Result<()> {
        let mut stores = self.stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        if stores.contains_key(alias) {
            return Err(FederationError::ConfigError(
                format!("Store with alias '{}' already registered", alias)
            ));
        }
        
        stores.insert(alias.to_string(), StoreEntry {
            store: store.clone(),
            store_type: StoreType::Vector,
        });
        
        let mut vector_stores = self.vector_stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        vector_stores.insert(alias.to_string(), store.clone());
        
        // Also register with transaction coordinator for 2PC support
        self.tx_coordinator.register_store(alias, store)?;
        
        Ok(())
    }
    
    /// Unregister a store
    pub fn unregister(&self, alias: &str) -> Result<()> {
        let mut stores = self.stores.write()
            .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
        
        if let Some(entry) = stores.remove(alias) {
            match entry.store_type {
                StoreType::Graph => {
                    let mut graph_stores = self.graph_stores.write()
                        .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
                    graph_stores.remove(alias);
                }
                StoreType::Document => {
                    let mut doc_stores = self.doc_stores.write()
                        .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
                    doc_stores.remove(alias);
                    // Also try removing from concurrent doc stores
                    let mut concurrent_doc_stores = self.concurrent_doc_stores.write()
                        .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
                    concurrent_doc_stores.remove(alias);
                }
                StoreType::Vector => {
                    let mut vector_stores = self.vector_stores.write()
                        .map_err(|_| FederationError::LockError("Failed to acquire write lock".to_string()))?;
                    vector_stores.remove(alias);
                }
            }
        }
        
        Ok(())
    }
    
    // ===== Type-Safe Store Access =====
    
    /// Get a store by alias (generic)
    pub fn use_store(&self, alias: &str) -> Result<Arc<dyn Store>> {
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        stores.get(alias)
            .map(|e| e.store.clone())
            .ok_or_else(|| FederationError::StoreNotFound(alias.to_string()))
    }
    
    /// Get a graph store by alias (type-safe)
    pub fn use_graph(&self, alias: &str) -> Result<Arc<GraphStoreAdapter>> {
        let graph_stores = self.graph_stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        graph_stores.get(alias)
            .cloned()
            .ok_or_else(|| FederationError::StoreNotFound(alias.to_string()))
    }
    
    /// Get a document store by alias (type-safe)
    pub fn use_doc(&self, alias: &str) -> Result<Arc<DocumentStoreAdapter>> {
        let doc_stores = self.doc_stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        doc_stores.get(alias)
            .cloned()
            .ok_or_else(|| FederationError::StoreNotFound(alias.to_string()))
    }
    
    /// Get a vector store by alias (type-safe)
    pub fn use_vector(&self, alias: &str) -> Result<Arc<VectorStoreAdapter>> {
        let vector_stores = self.vector_stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        vector_stores.get(alias)
            .cloned()
            .ok_or_else(|| FederationError::StoreNotFound(alias.to_string()))
    }
    
    /// List all registered stores
    pub fn list_stores(&self) -> Result<Vec<(String, StoreType)>> {
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        Ok(stores.iter()
            .map(|(alias, entry)| (alias.clone(), entry.store_type))
            .collect())
    }
    
    // ===== Unified Query =====
    
    /// Execute a query across stores
    pub fn query(&self, store_alias: &str, q: &UniversalQuery) -> Result<QueryResult> {
        let start = std::time::Instant::now();
        
        let store = self.use_store(store_alias)?;
        let records = store.query(q)?;
        
        let stats = QueryStats {
            execution_time_ms: start.elapsed().as_millis() as u64,
            records_scanned: records.len(),
            records_returned: records.len(),
            join_strategy: None,
            stores_involved: vec![store_alias.to_string()],
        };
        
        let total = records.len();
        Ok(QueryResult {
            records,
            total_count: Some(total),
            has_more: false,
            stats,
        })
    }
    
    /// Execute a query with streaming results
    pub fn query_stream(&self, store_alias: &str, q: &UniversalQuery) -> Result<Box<dyn Iterator<Item = Result<Record>> + Send>> {
        let store = self.use_store(store_alias)?;
        store.query_iter(q)
    }
    
    /// Query multiple stores
    pub fn query_multi(&self, queries: Vec<(&str, &UniversalQuery)>) -> Result<Vec<QueryResult>> {
        queries.into_iter()
            .map(|(alias, q)| self.query(alias, q))
            .collect()
    }
    
    /// Execute a cross-database query using the query planner and executor.
    ///
    /// This is the main entry point for federated queries that span multiple
    /// heterogeneous stores (e.g. Graph + Document).
    ///
    /// # Example
    /// ```rust,ignore
    /// let query = CrossDbQuery {
    ///     targets: vec![
    ///         QueryTarget::graph("social", "person"),
    ///         QueryTarget::document("knowledge", "article"),
    ///     ],
    ///     joins: vec![JoinClause {
    ///         target: QueryTarget::document("knowledge", "article"),
    ///         on: "person.id = article.author_id".to_string(),
    ///         join_type: JoinType::Inner,
    ///     }],
    ///     filters: vec![Filter::Eq("person.name".into(), json!("Alice"))],
    ///     select: None,
    ///     order_by: None,
    ///     limit: Some(20),
    ///     offset: None,
    ///     streaming: false,
    ///     batch_size: None,
    /// };
    /// let result = federation.cross_query(&query)?;
    /// ```
    pub fn cross_query(&self, query: &CrossDbQuery) -> Result<QueryResult> {
        let start = std::time::Instant::now();
        
        // 1. Plan the query
        let plan = self.planner.plan_cross_db(query)?;
        
        // 2. Build executor with all involved stores
        let mut executor = QueryExecutor::new(ExecutorConfig {
            timeout_ms: self.config.query_timeout_ms,
            ..ExecutorConfig::default()
        });
        
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        for alias in &plan.stores_involved {
            if let Some(entry) = stores.get(alias) {
                executor.register_store(alias, entry.store.clone());
            } else {
                return Err(FederationError::StoreNotFound(alias.clone()));
            }
        }
        
        // 3. Execute the plan
        let mut result = executor.execute(&plan)?;
        
        // 4. Patch statistics with overall timing
        result.stats.execution_time_ms = start.elapsed().as_millis() as u64;
        result.stats.stores_involved = plan.stores_involved.clone();
        
        Ok(result)
    }
    
    // ===== Unified Write =====
    
    /// Write to a specific store
    pub fn write(&self, store_alias: &str, op: WriteOp) -> Result<super::types::WriteResult> {
        let store = self.use_store(store_alias)?;
        store.write(op)
    }
    
    /// Batch write to a specific store
    pub fn batch_write(&self, store_alias: &str, ops: Vec<WriteOp>) -> Result<BatchWriteResult> {
        let store = self.use_store(store_alias)?;
        store.batch_write(ops)
    }
    
    // ===== Transactions =====
    
    /// Begin a cross-store transaction
    pub fn begin_transaction(&self, store_aliases: Vec<&str>) -> Result<FederatedTransaction> {
        let tx_id = self.tx_coordinator.begin()?;
        let mut stores = Vec::new();
        
        for alias in store_aliases {
            let store = self.use_store(alias)?;
            stores.push((alias.to_string(), store));
        }
        
        Ok(FederatedTransaction {
            id: tx_id,
            stores,
            coordinator: self.tx_coordinator.clone(),
        })
    }
    
    // ===== Index Management =====
    
    /// Rebuild cross-database index
    pub fn rebuild_index(&self) -> Result<()> {
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        self.index.clear()?;
        
        for (alias, entry) in stores.iter() {
            self.index.index_store(alias, &*entry.store)?;
        }
        
        Ok(())
    }
    
    /// Get index statistics
    pub fn index_stats(&self) -> super::cross_db_index::IndexStats {
        self.index.stats()
    }
    
    // ===== Lifecycle =====
    
    /// Sync all stores to disk
    pub fn sync_all(&self) -> Result<()> {
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        for (_, entry) in stores.iter() {
            entry.store.sync()?;
        }
        
        Ok(())
    }
    
    /// Close all stores
    pub fn close(&self) -> Result<()> {
        let stores = self.stores.read()
            .map_err(|_| FederationError::LockError("Failed to acquire read lock".to_string()))?;
        
        for (_, entry) in stores.iter() {
            entry.store.close()?;
        }
        
        Ok(())
    }
    
    // ===== Authenticated Operations (Layer 1 RBAC + Layer 2 ABAC) =====
    
    /// Get a reference to the permission checker (for external configuration)
    pub fn permission_checker(&self) -> &PermissionChecker {
        &self.permission_checker
    }
    
    /// Write with authentication and authorization checks
    ///
    /// Performs two-layer permission checking:
    /// - **Layer 1 (RBAC)**: Checks whether the caller's roles grant `StoreWrite` permission
    /// - **Layer 2 (ABAC)**: For Update/Delete operations, checks record-level ownership
    ///   and protection levels against the existing record's `OwnershipMeta`
    ///
    /// For Create operations, this method automatically sets `OwnershipMeta` using the
    /// caller's identity (from `ctx.caller`) if the WriteOp carries an `owner_id`.
    pub fn write_with_auth(
        &self,
        alias: &str,
        op: WriteOp,
        ctx: &AccessContext,
    ) -> Result<super::types::WriteResult> {
        // Layer 1: RBAC check
        self.permission_checker.check(ctx, Permission::StoreWrite)?;
        
        // Layer 2: Record-level permission check for Update/Delete
        if op.is_update() || op.is_delete() {
            let record_id = op.record_id().to_string();
            let store = self.use_store(alias)?;
            
            if let Some(record) = store.get(&record_id)? {
                if let Some(ownership) = record.ownership() {
                    if op.is_update() {
                        if !RecordPermissionChecker::can_edit(ctx, ownership) {
                            return Err(FederationError::PermissionDenied {
                                caller: ctx.caller.clone(),
                                required: format!(
                                    "record:{}:edit (owner: {})",
                                    record_id, ownership.owner_id
                                ),
                            });
                        }
                    } else if op.is_delete() {
                        if !RecordPermissionChecker::can_delete(ctx, ownership) {
                            return Err(FederationError::PermissionDenied {
                                caller: ctx.caller.clone(),
                                required: format!(
                                    "record:{}:delete (owner: {})",
                                    record_id, ownership.owner_id
                                ),
                            });
                        }
                    }
                }
                // If ownership is None, skip ABAC check (backward compatible)
            }
            // If record doesn't exist yet, skip ABAC check (will fail at write level)
        }
        
        // Execute the write
        self.write(alias, op)
    }
    
    /// Query with authentication and visibility filtering
    ///
    /// Performs two-layer permission checking:
    /// - **Layer 1 (RBAC)**: Checks whether the caller has query permissions
    /// - **Layer 2 (ABAC)**: Filters out records that the caller cannot read
    ///   based on ownership visibility rules (Private/Shared/RoleRestricted)
    ///
    /// Records without `OwnershipMeta` are treated as publicly visible (backward compatible).
    pub fn query_with_auth(
        &self,
        store_alias: &str,
        q: &UniversalQuery,
        ctx: &AccessContext,
    ) -> Result<QueryResult> {
        // Layer 1: RBAC check
        self.permission_checker.check(ctx, Permission::SingleStoreQuery)?;
        
        // Execute query
        let mut result = self.query(store_alias, q)?;
        
        // Layer 2: Filter out records that the caller cannot read
        let original_count = result.records.len();
        result.records.retain(|record| {
            match record.ownership() {
                Some(ownership) => RecordPermissionChecker::can_read(ctx, ownership),
                None => true, // No ownership → default visible (backward compatible)
            }
        });
        
        // Update statistics
        let _filtered_count = original_count - result.records.len();
        result.stats.records_returned = result.records.len();
        result.total_count = Some(result.records.len());
        
        Ok(result)
    }
    
    /// Write system-level protected data (internal use only)
    ///
    /// This method is intended for:
    /// - Database initialization (seed data)
    /// - System upgrade/migration scripts
    /// - Internal service calls
    ///
    /// It automatically:
    /// - Sets `OwnershipMeta` with `SYSTEM_OWNER_ID` as owner
    /// - Applies the specified `ProtectionLevel`
    /// - Uses `system_override` to bypass all permission checks
    pub fn write_system_data(
        &self,
        alias: &str,
        mut op: WriteOp,
        protection: ProtectionLevel,
    ) -> Result<super::types::WriteResult> {
        // Inject system owner_id and protection into the WriteOp
        Self::inject_system_owner(&mut op);
        Self::inject_protection(&mut op, &protection);
        
        // Execute the write
        let result = self.write(alias, op)?;
        
        Ok(result)
    }
    
    /// Admin: downgrade a record's protection level
    ///
    /// Use case: An Admin needs to edit/delete an Immutable record.
    /// Flow: downgrade_protection(Immutable → None) → edit/delete → done
    ///
    /// Rules:
    /// - Only Admin or `system_override` can change protection levels
    /// - `System` level can only be modified with `system_override`
    /// - Cannot upgrade to `System` via API (requires `system_override`)
    pub fn downgrade_protection(
        &self,
        alias: &str,
        record_id: &str,
        new_level: ProtectionLevel,
        ctx: &AccessContext,
    ) -> Result<()> {
        // Get the current record
        let store = self.use_store(alias)?;
        let record = store
            .get(record_id)?
            .ok_or_else(|| FederationError::RecordNotFound {
                id: record_id.to_string(),
            })?;
        
        let ownership = record.ownership().ok_or_else(|| FederationError::NoOwnership {
            id: record_id.to_string(),
        })?;
        
        // Check if the protection level change is allowed
        match RecordPermissionChecker::check_protection_change(ctx, &ownership.protection, &new_level) {
            ProtectionCheckResult::Denied { reason } => {
                Err(FederationError::ProtectionError { reason })
            }
            ProtectionCheckResult::Continue => {
                // Protection level change is allowed — persist it
                let mut updated_ownership = ownership.clone();
                updated_ownership.protection = new_level;
                store.update_ownership(record_id, updated_ownership)?;
                Ok(())
            }
        }
    }
    
    /// Inject SYSTEM_OWNER_ID into a WriteOp's owner_id field
    fn inject_system_owner(op: &mut WriteOp) {
        match op {
            WriteOp::CreateEntity { owner_id, visibility, .. } => {
                *owner_id = Some(SYSTEM_OWNER_ID.to_string());
                if visibility.is_none() {
                    *visibility = Some(Visibility::Public);
                }
            }
            WriteOp::CreateDocument { owner_id, visibility, .. } => {
                *owner_id = Some(SYSTEM_OWNER_ID.to_string());
                if visibility.is_none() {
                    *visibility = Some(Visibility::Public);
                }
            }
            WriteOp::UpsertVector { owner_id, visibility, .. } => {
                *owner_id = Some(SYSTEM_OWNER_ID.to_string());
                if visibility.is_none() {
                    *visibility = Some(Visibility::Public);
                }
            }
            WriteOp::CreateRelation { owner_id, .. } => {
                *owner_id = Some(SYSTEM_OWNER_ID.to_string());
            }
            _ => {} // Update/Delete ops don't have owner_id to set
        }
    }
    
    /// Inject protection level into WriteOp's content metadata
    /// 
    /// For CreateDocument ops, stores the protection level as `__protection__`
    /// in the content JSON so the adapter can extract it when building the MemoryNode.
    fn inject_protection(op: &mut WriteOp, protection: &ProtectionLevel) {
        match op {
            WriteOp::CreateDocument { content, .. } => {
                if let Some(obj) = content.as_object_mut() {
                    obj.insert(
                        "__protection__".to_string(),
                        serde_json::to_value(protection).unwrap_or(serde_json::Value::Null),
                    );
                }
            }
            _ => {} // Other store types will be extended in future iterations
        }
    }
}

/// Federated transaction across multiple stores
pub struct FederatedTransaction {
    id: u64,
    stores: Vec<(String, Arc<dyn Store>)>,
    coordinator: Arc<TransactionCoordinator>,
}

impl FederatedTransaction {
    /// Execute a write operation within the transaction
    pub fn write(&mut self, store_alias: &str, op: WriteOp) -> Result<()> {
        self.coordinator.add_operation(self.id, store_alias, op)
    }
    
    /// Commit the transaction (2PC)
    pub fn commit(self) -> Result<()> {
        self.coordinator.commit(self.id)
    }
    
    /// Rollback the transaction
    pub fn rollback(self) -> Result<()> {
        self.coordinator.rollback(self.id)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_federation_creation() {
        let fed = Federation::with_defaults();
        let stores = fed.list_stores().unwrap();
        assert!(stores.is_empty());
    }
    
    #[test]
    fn test_store_registration() {
        let fed = Federation::with_defaults();
        
        let temp_dir = std::env::temp_dir().join("aifastdb_test_fed_reg");
        let _ = std::fs::create_dir_all(&temp_dir);
        let doc_path = temp_dir.join("docs.jsonl");
        
        let doc_store = DocumentStoreAdapter::open(
            "test_docs".to_string(),
            doc_path.to_str().unwrap(),
        ).unwrap();
        let doc_store = Arc::new(doc_store);
        
        fed.register_document("test_docs", doc_store).unwrap();
        
        let stores = fed.list_stores().unwrap();
        assert_eq!(stores.len(), 1);
        assert_eq!(stores[0].0, "test_docs");
        assert_eq!(stores[0].1, StoreType::Document);
        
        // Cleanup
        let _ = std::fs::remove_dir_all(&temp_dir);
    }
}
