//! Python bindings for Federation & EnhancedDocumentStore.
//!
//! Provides:
//! - `DocumentStore` class wrapping `EnhancedDocumentStore` (T27.1)
//! - `Federation` class wrapping `Federation` (T27.2 – T27.7)
//!   - Store registration (T27.3)
//!   - Unified query (T27.4)
//!   - Unified write (T27.5)
//!   - Transaction support (T27.6)
//!   - Permission-checked operations (T27.7)

use std::collections::HashMap;
use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::exceptions::{PyValueError, PyRuntimeError};

use aifastdb_core::storage::document_enhanced::{EnhancedDocumentStore, EnhancedStoreConfig};
use aifastdb_core::types::{ContentType, MemoryNode};
use aifastdb_core::federation::{
    Federation, FederationConfig,
    GraphStoreAdapter, DocumentStoreAdapter, VectorStoreAdapter, ConcurrentDocStoreAdapter,
    StoreType as CoreStoreType,
    WriteOp, UniversalQuery, CrossDbQuery, Filter, SortOrder,
    AccessContext, RecordPermissionChecker,
};
use aifastdb_core::federation::types::{Record, QueryResult};
use aifastdb_core::social::{SocialStoreV2, SocialStoreV2Config};
use aifastdb_core::vector::concurrent::{ConcurrentHNSWConfig, ConcurrentHNSWIndex};
use aifastdb_core::storage::concurrent_store::{ConcurrentDocumentStore, ConcurrentStoreConfig};
use aifastdb_core::types::config::HNSWConfig;

use crate::errors::federation_error_to_pyerr;
use crate::types::{json_value_to_py, py_to_json_value};

// ============================================================================
// Helper: PyDict → serde_json::Value
// ============================================================================

fn pydict_to_json(dict: &Bound<'_, PyDict>) -> PyResult<serde_json::Value> {
    let mut map = serde_json::Map::new();
    for (key, val) in dict.iter() {
        let key_str: String = key.extract()?;
        map.insert(key_str, py_to_json_value(&val)?);
    }
    Ok(serde_json::Value::Object(map))
}

fn json_to_py(py: Python<'_>, val: &serde_json::Value) -> PyResult<PyObject> {
    json_value_to_py(py, val)
}

// ============================================================================
// T27.1: DocumentStore class (wraps EnhancedDocumentStore directly)
// ============================================================================

/// DocumentStore — Python interface for the enhanced JSONL document store.
///
/// Usage:
///     with DocumentStore("/path/to/data.jsonl") as store:
///         store.put("doc1", "Hello", doc_type="text", tags=["greeting"])
///         doc = store.get("doc1")
#[pyclass(name = "DocumentStore")]
pub struct PyDocumentStore {
    store: Arc<EnhancedDocumentStore>,
    path: String,
}

#[pymethods]
impl PyDocumentStore {
    /// Create or open a DocumentStore.
    ///
    /// Args:
    ///     path: Path to the JSONL data file.
    ///     flush_interval_ms: Flush interval in ms (default 1000).
    ///     enable_indexing: Enable in-memory indexing (default True).
    ///     enable_file_lock: Enable file locking for multi-process safety (default True).
    #[new]
    #[pyo3(signature = (path, *, flush_interval_ms=1000, enable_indexing=true, enable_file_lock=true))]
    fn new(path: String, flush_interval_ms: u64, enable_indexing: bool, enable_file_lock: bool) -> PyResult<Self> {
        let config = EnhancedStoreConfig::new()
            .with_flush_interval_ms(flush_interval_ms)
            .with_indexing(enable_indexing)
            .with_file_lock(enable_file_lock);

        let store = EnhancedDocumentStore::open(&path, config)
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to open document store: {}", e)))?;
        Ok(Self {
            store: Arc::new(store),
            path,
        })
    }

    /// Context manager entry.
    fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Context manager exit — flush.
    #[pyo3(signature = (_exc_type=None, _exc_val=None, _exc_tb=None))]
    fn __exit__(
        &self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<bool> {
        let _ = self.store.flush();
        Ok(false)
    }

    #[getter]
    fn path(&self) -> &str {
        &self.path
    }

    fn __repr__(&self) -> String {
        format!("DocumentStore(path='{}', len={})", self.path, self.store.len())
    }

    // ---- CRUD ----

    /// Store a document.
    ///
    /// Args:
    ///     id: Document ID.
    ///     content: Text content.
    ///     doc_type: Content type string ("text", "code", "document", etc.). Default "text".
    ///     tags: Optional list of tag strings.
    ///     metadata: Optional dict of metadata.
    ///     importance: Importance score 0.0–1.0. Default 0.5.
    #[pyo3(signature = (id, content, *, doc_type="text", tags=None, metadata=None, importance=0.5))]
    fn put(
        &self,
        id: &str,
        content: &str,
        doc_type: &str,
        tags: Option<Vec<String>>,
        metadata: Option<&Bound<'_, PyDict>>,
        importance: f64,
    ) -> PyResult<()> {
        let content_type = match doc_type {
            "text" => ContentType::Text,
            "image" => ContentType::Image,
            "code" => ContentType::Code,
            "conversation" => ContentType::Conversation,
            "document" => ContentType::Document,
            _ => ContentType::Other,
        };

        let meta: HashMap<String, serde_json::Value> = match metadata {
            Some(dict) => {
                let json = pydict_to_json(dict)?;
                serde_json::from_value(json).unwrap_or_default()
            }
            None => HashMap::new(),
        };

        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as i64;

        let node = MemoryNode {
            id: id.to_string(),
            content: content.to_string(),
            embedding: None,
            content_type,
            parent_id: None,
            children: vec![],
            depth: 0,
            metadata: meta,
            tags: tags.unwrap_or_default(),
            created_at: now,
            last_accessed_at: now,
            access_count: 0,
            importance,
            decay_rate: 0.0001,
            entity_ids: vec![],
            relation_ids: vec![],
            owner_id: None,
            visibility: Default::default(),
            protection: Default::default(),
            collaborators: vec![],
            semantic_axes: None,
        };

        self.store
            .put(node)
            .map_err(|e| PyRuntimeError::new_err(format!("Put failed: {}", e)))
    }

    /// Get a document by ID.
    ///
    /// Returns dict with "id", "content", "content_type", "tags", "metadata",
    /// "importance", "created_at", "last_accessed_at", or None if not found.
    fn get(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        match self.store.get(id) {
            Some(node) => Ok(Some(node_to_py(py, &node)?)),
            None => Ok(None),
        }
    }

    /// Check if a document exists.
    fn has(&self, id: &str) -> bool {
        self.store.contains(id)
    }

    /// Delete a document by ID.
    ///
    /// Returns the deleted document dict, or None if not found.
    fn delete(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        match self.store.delete(id) {
            Some(node) => Ok(Some(node_to_py(py, &node)?)),
            None => Ok(None),
        }
    }

    /// Find documents by tag.
    ///
    /// Returns list of document dicts.
    fn find_by_tag(&self, py: Python<'_>, tag: &str) -> PyResult<Vec<PyObject>> {
        let nodes = self.store.find_by_tag(tag);
        nodes_to_py_vec(py, &nodes)
    }

    /// Find documents by content type.
    ///
    /// Args:
    ///     doc_type: Content type string ("text", "code", "document", etc.).
    ///
    /// Returns list of document dicts.
    fn find_by_type(&self, py: Python<'_>, doc_type: &str) -> PyResult<Vec<PyObject>> {
        let content_type = match doc_type {
            "text" => ContentType::Text,
            "image" => ContentType::Image,
            "code" => ContentType::Code,
            "conversation" => ContentType::Conversation,
            "document" => ContentType::Document,
            _ => ContentType::Other,
        };
        let nodes = self.store.find_by_type(content_type);
        nodes_to_py_vec(py, &nodes)
    }

    /// List all document IDs.
    fn ids(&self) -> Vec<String> {
        self.store.ids()
    }

    /// Get the total number of documents.
    fn __len__(&self) -> usize {
        self.store.len()
    }

    /// Check if the store is empty.
    fn is_empty(&self) -> bool {
        self.store.is_empty()
    }

    /// Get all tags used across all documents.
    fn all_tags(&self) -> Vec<String> {
        self.store.all_tags()
    }

    /// Flush pending writes to disk.
    fn flush(&self) -> PyResult<()> {
        self.store
            .flush()
            .map_err(|e| PyRuntimeError::new_err(format!("Flush failed: {}", e)))
    }

    /// Compact the data file (remove deleted entries).
    fn compact(&self) -> PyResult<()> {
        self.store
            .compact()
            .map_err(|e| PyRuntimeError::new_err(format!("Compact failed: {}", e)))
    }

    /// Find documents by time range.
    ///
    /// Args:
    ///     start: Start timestamp (Unix ms).
    ///     end: End timestamp (Unix ms).
    ///
    /// Returns list of document dicts.
    fn find_by_time_range(&self, py: Python<'_>, start: i64, end: i64) -> PyResult<Vec<PyObject>> {
        let nodes = self.store.find_by_time_range(start, end);
        nodes_to_py_vec(py, &nodes)
    }

    /// Find documents by minimum importance.
    fn find_by_min_importance(&self, py: Python<'_>, min_importance: f64) -> PyResult<Vec<PyObject>> {
        let nodes = self.store.find_by_min_importance(min_importance);
        nodes_to_py_vec(py, &nodes)
    }

    /// Get all documents.
    fn all(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let nodes = self.store.all();
        nodes_to_py_vec(py, &nodes)
    }
}

// ============================================================================
// T27.2–T27.7: Federation class
// ============================================================================

/// Federation — Python interface for heterogeneous storage federation.
///
/// Manages multiple stores (graph, document, vector) through a unified interface.
///
/// Usage:
///     with Federation() as fed:
///         fed.register_graph("social", "/path/to/social")
///         fed.register_document("docs", "/path/to/docs.jsonl")
///         result = fed.query("social", {"target": "person"})
#[pyclass(name = "Federation")]
pub struct PyFederation {
    inner: Federation,
    /// Optional tokio runtime for async store adapters (vector, concurrent doc).
    /// Created lazily when register_vector() or register_concurrent_document() is called.
    runtime: Option<tokio::runtime::Runtime>,
}

#[pymethods]
impl PyFederation {
    /// Create a new Federation instance.
    ///
    /// Args:
    ///     enable_permissions: Enable RBAC permission checking. Default False.
    ///     query_timeout_ms: Default query timeout in ms. Default 30000.
    #[new]
    #[pyo3(signature = (*, enable_permissions=false, query_timeout_ms=30000))]
    fn new(enable_permissions: bool, query_timeout_ms: u64) -> PyResult<Self> {
        let config = FederationConfig {
            enable_permissions,
            query_timeout_ms,
            ..FederationConfig::default()
        };
        Ok(Self {
            inner: Federation::new(config),
            runtime: None,
        })
    }

    /// Context manager entry.
    fn __enter__(slf: Py<Self>) -> Py<Self> {
        slf
    }

    /// Context manager exit — sync + close.
    ///
    /// The runtime context was already entered permanently in `ensure_runtime()`,
    /// so `Handle::current()` is available for the adapters' `block_on_async()`.
    #[pyo3(signature = (_exc_type=None, _exc_val=None, _exc_tb=None))]
    fn __exit__(
        &self,
        _exc_type: Option<&Bound<'_, PyAny>>,
        _exc_val: Option<&Bound<'_, PyAny>>,
        _exc_tb: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<bool> {
        let _ = self.inner.sync_all();
        let _ = self.inner.close();
        Ok(false)
    }

    fn __repr__(&self) -> PyResult<String> {
        let stores = self
            .inner
            .list_stores()
            .map_err(federation_error_to_pyerr)?;
        Ok(format!(
            "Federation(stores={})",
            stores.len()
        ))
    }

    // ========================================================================
    // T27.3: Store Registration
    // ========================================================================

    /// Register a graph store (SocialGraphV2).
    ///
    /// Args:
    ///     alias: Store alias name.
    ///     path: Path to the data directory.
    ///     shard_count: Number of shards (0 = auto). Default 0.
    ///     wal_enabled: Enable WAL. Default True.
    #[pyo3(signature = (alias, path, *, shard_count=0, wal_enabled=true))]
    fn register_graph(
        &self,
        alias: &str,
        path: &str,
        shard_count: usize,
        wal_enabled: bool,
    ) -> PyResult<()> {
        let mut config = SocialStoreV2Config::new(path);
        if shard_count > 0 {
            config = config.with_shard_count(shard_count);
        }
        config = config.with_wal(wal_enabled);

        let store = SocialStoreV2::new(config);
        let adapter = GraphStoreAdapter::new(alias.to_string(), Arc::new(store));
        self.inner
            .register_graph(alias, Arc::new(adapter))
            .map_err(federation_error_to_pyerr)
    }

    /// Register a document store (EnhancedDocumentStore).
    ///
    /// Args:
    ///     alias: Store alias name.
    ///     path: Path to the JSONL data file.
    #[pyo3(signature = (alias, path, *, enable_file_lock=true))]
    fn register_document(&self, alias: &str, path: &str, enable_file_lock: bool) -> PyResult<()> {
        let config = EnhancedStoreConfig::new().with_file_lock(enable_file_lock);
        let adapter = DocumentStoreAdapter::open_with_config(alias.to_string(), path, config)
            .map_err(federation_error_to_pyerr)?;
        self.inner
            .register_document(alias, Arc::new(adapter))
            .map_err(federation_error_to_pyerr)
    }

    /// Register a vector store (ConcurrentHNSWIndex).
    ///
    /// Args:
    ///     alias: Store alias name.
    ///     dimension: Vector dimension (required).
    ///     shard_count: Number of shards (0 = auto).
    ///     mode: "default", "high_throughput", or "low_latency".
    ///     m: HNSW M parameter.
    ///     ef_construction: HNSW ef_construction parameter.
    ///     ef_search: HNSW ef_search parameter.
    ///     max_elements: Maximum number of vectors.
    #[pyo3(signature = (alias, *, dimension=128, shard_count=0, mode="default",
                        m=16, ef_construction=100, ef_search=50, max_elements=100000))]
    fn register_vector(
        &mut self,
        alias: &str,
        dimension: usize,
        shard_count: usize,
        mode: &str,
        m: usize,
        ef_construction: usize,
        ef_search: usize,
        max_elements: usize,
    ) -> PyResult<()> {
        self.ensure_runtime()?;

        let hnsw = HNSWConfig {
            dimension,
            m,
            m_max0: m * 2,
            ef_construction,
            ef_search,
            max_elements,
            seed: None,
        };

        let config = match mode {
            "high_throughput" => ConcurrentHNSWConfig::high_throughput(hnsw),
            "low_latency" => ConcurrentHNSWConfig::low_latency(hnsw),
            _ => ConcurrentHNSWConfig::with_hnsw(hnsw),
        };

        let config = if shard_count > 0 { config.with_shards(shard_count) } else { config };

        let adapter = VectorStoreAdapter::with_config(alias.to_string(), config);
        self.inner
            .register_vector(alias, Arc::new(adapter))
            .map_err(federation_error_to_pyerr)
    }

    /// Register a high-concurrency document store (ConcurrentDocumentStore).
    ///
    /// Args:
    ///     alias: Store alias name.
    ///     path: Directory path for the concurrent store data.
    ///     shard_count: Number of write shards (0 = auto).
    ///     enable_wal: Whether to enable Write-Ahead Logging.
    ///     mode: "default", "high_throughput", or "low_latency".
    #[pyo3(signature = (alias, path, *, shard_count=0, enable_wal=true, mode="default"))]
    fn register_concurrent_document(
        &mut self,
        alias: &str,
        path: &str,
        shard_count: usize,
        enable_wal: bool,
        mode: &str,
    ) -> PyResult<()> {
        self.ensure_runtime()?;

        let mut config = ConcurrentStoreConfig::new(path);
        if shard_count > 0 {
            config = config.with_shards(shard_count);
        }
        if !enable_wal {
            config = config.without_wal();
        }
        config = match mode {
            "high_throughput" => config.high_throughput(),
            "low_latency" => config.low_latency(),
            _ => config,
        };

        let rt = self.runtime.as_ref().unwrap();
        let store = rt.block_on(ConcurrentDocumentStore::new(config))
            .map_err(|e| PyRuntimeError::new_err(format!("Failed to create concurrent store: {}", e)))?;

        let adapter = ConcurrentDocStoreAdapter::new(alias.to_string(), Arc::new(store));
        self.inner
            .register_concurrent_document(alias, Arc::new(adapter))
            .map_err(federation_error_to_pyerr)
    }

    /// Unregister a store by alias.
    fn unregister(&self, alias: &str) -> PyResult<()> {
        self.inner
            .unregister(alias)
            .map_err(federation_error_to_pyerr)
    }

    /// List all registered stores.
    ///
    /// Returns list of (alias, store_type) tuples.
    fn list_stores(&self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
        let stores = self
            .inner
            .list_stores()
            .map_err(federation_error_to_pyerr)?;

        let mut results = Vec::with_capacity(stores.len());
        for (alias, store_type) in stores {
            let type_str = match store_type {
                CoreStoreType::Graph => "graph",
                CoreStoreType::Document => "document",
                CoreStoreType::Vector => "vector",
            };
            let tuple = (alias, type_str.to_string()).into_py(py);
            results.push(tuple);
        }
        Ok(results)
    }

    // ========================================================================
    // T27.4: Unified Query
    // ========================================================================

    /// Execute a query on a single store.
    ///
    /// Args:
    ///     store_alias: Alias of the target store.
    ///     query: Query dict with:
    ///         - "target": target type (e.g. "person", "text")
    ///         - "filters": list of filter dicts (optional)
    ///         - "order_by": [field, "asc"|"desc"] (optional)
    ///         - "limit": int (optional)
    ///         - "offset": int (optional)
    ///
    /// Returns dict with "records", "total_count", "has_more", "stats".
    fn query(&self, py: Python<'_>, store_alias: &str, query: &Bound<'_, PyDict>) -> PyResult<PyObject> {
        let q = parse_universal_query(query)?;
        let result = self
            .inner
            .query(store_alias, &q)
            .map_err(federation_error_to_pyerr)?;
        query_result_to_py(py, &result)
    }

    /// Execute a cross-database query.
    ///
    /// Args:
    ///     query: CrossDbQuery dict with:
    ///         - "targets": list of target dicts
    ///         - "joins": list of join clause dicts
    ///         - "filters": list of filter dicts
    ///         - "limit", "offset": pagination
    ///
    /// Returns dict with "records", "total_count", "has_more", "stats".
    fn cross_query(&self, py: Python<'_>, query: &Bound<'_, PyDict>) -> PyResult<PyObject> {
        let json = pydict_to_json(query)?;
        let cross_q: CrossDbQuery = serde_json::from_value(json)
            .map_err(|e| PyValueError::new_err(format!("Invalid cross query: {}", e)))?;

        let result = self
            .inner
            .cross_query(&cross_q)
            .map_err(federation_error_to_pyerr)?;
        query_result_to_py(py, &result)
    }

    // ========================================================================
    // T27.5: Unified Write
    // ========================================================================

    /// Write to a specific store.
    ///
    /// Args:
    ///     store_alias: Alias of the target store.
    ///     operation: Write operation dict. Examples:
    ///         - {"op": "CreateDocument", "id": "d1", "doc_type": "text",
    ///            "content": {"text": "hello"}, "tags": ["greeting"]}
    ///         - {"op": "CreateEntity", "id": "e1", "entity_type": "person",
    ///            "name": "Alice", "attributes": {}}
    ///         - {"op": "DeleteDocument", "id": "d1"}
    ///
    /// Returns str describing the result (e.g. "Created(d1)").
    fn write(&self, store_alias: &str, operation: &Bound<'_, PyDict>) -> PyResult<String> {
        let op = parse_write_op(operation)?;
        let result = self
            .inner
            .write(store_alias, op)
            .map_err(federation_error_to_pyerr)?;
        Ok(format!("{:?}", result))
    }

    /// Batch write to a specific store.
    ///
    /// Args:
    ///     store_alias: Alias of the target store.
    ///     operations: List of write operation dicts.
    ///
    /// Returns dict with "success_count", "failed_count", "execution_time_ms".
    fn batch_write(
        &self,
        py: Python<'_>,
        store_alias: &str,
        operations: Vec<Bound<'_, PyDict>>,
    ) -> PyResult<PyObject> {
        let mut ops = Vec::with_capacity(operations.len());
        for op_dict in &operations {
            ops.push(parse_write_op(op_dict)?);
        }

        let result = self
            .inner
            .batch_write(store_alias, ops)
            .map_err(federation_error_to_pyerr)?;

        let val = serde_json::to_value(&result)
            .map_err(|e| PyRuntimeError::new_err(format!("Serialization error: {}", e)))?;
        json_to_py(py, &val)
    }

    // ========================================================================
    // T27.6: Transaction Support
    // ========================================================================

    /// Begin a transaction across one or more stores.
    ///
    /// Args:
    ///     store_aliases: List of store alias strings to include in the transaction.
    ///
    /// Returns a FederatedTransaction object.
    fn begin_transaction(&self, store_aliases: Vec<String>) -> PyResult<PyFederatedTransaction> {
        let alias_refs: Vec<&str> = store_aliases.iter().map(|s| s.as_str()).collect();
        let tx = self
            .inner
            .begin_transaction(alias_refs)
            .map_err(federation_error_to_pyerr)?;
        Ok(PyFederatedTransaction {
            inner: Some(tx),
        })
    }

    // ========================================================================
    // T27.7: Permission-Checked Operations
    // ========================================================================

    /// Write with authentication and authorization checks.
    ///
    /// Args:
    ///     store_alias: Target store alias.
    ///     operation: Write operation dict.
    ///     caller: Caller identifier (user ID).
    ///     roles: List of role strings (e.g. ["admin"], ["writer"]).
    ///
    /// Returns str describing the result.
    fn write_with_auth(
        &self,
        store_alias: &str,
        operation: &Bound<'_, PyDict>,
        caller: &str,
        roles: Vec<String>,
    ) -> PyResult<String> {
        let op = parse_write_op(operation)?;
        let ctx = AccessContext::new(caller, roles);
        let result = self
            .inner
            .write_with_auth(store_alias, op, &ctx)
            .map_err(federation_error_to_pyerr)?;
        Ok(format!("{:?}", result))
    }

    /// Query with authentication and visibility filtering.
    ///
    /// Args:
    ///     store_alias: Target store alias.
    ///     query: Query dict (same format as `query()`).
    ///     caller: Caller identifier (user ID).
    ///     roles: List of role strings.
    ///
    /// Returns dict with "records", "total_count", "has_more", "stats".
    fn query_with_auth(
        &self,
        py: Python<'_>,
        store_alias: &str,
        query: &Bound<'_, PyDict>,
        caller: &str,
        roles: Vec<String>,
    ) -> PyResult<PyObject> {
        let q = parse_universal_query(query)?;
        let ctx = AccessContext::new(caller, roles);
        let result = self
            .inner
            .query_with_auth(store_alias, &q, &ctx)
            .map_err(federation_error_to_pyerr)?;
        query_result_to_py(py, &result)
    }

    /// Get available actions for a user on a specific record.
    ///
    /// Args:
    ///     store_alias: Target store alias.
    ///     record_id: Record ID.
    ///     caller: Caller identifier.
    ///     roles: List of role strings.
    ///
    /// Returns dict with "can_read", "can_edit", "can_delete", "can_share",
    /// "can_change_protection", "is_owner", "is_system_data", "protection".
    fn get_record_actions(
        &self,
        py: Python<'_>,
        store_alias: &str,
        record_id: &str,
        caller: &str,
        roles: Vec<String>,
    ) -> PyResult<PyObject> {
        let store = self
            .inner
            .use_store(store_alias)
            .map_err(federation_error_to_pyerr)?;
        let record = store
            .get(record_id)
            .map_err(federation_error_to_pyerr)?
            .ok_or_else(|| {
                PyValueError::new_err(format!("Record not found: {}", record_id))
            })?;

        let ctx = AccessContext::new(caller, roles);

        match record.ownership() {
            Some(ownership) => {
                let actions = RecordPermissionChecker::get_available_actions(&ctx, ownership);
                let val = serde_json::to_value(&actions)
                    .map_err(|e| PyRuntimeError::new_err(format!("Serialization error: {}", e)))?;
                json_to_py(py, &val)
            }
            None => {
                // No ownership metadata → return default (full access)
                let default_actions = serde_json::json!({
                    "can_read": true,
                    "can_edit": true,
                    "can_delete": true,
                    "can_share": false,
                    "can_change_protection": false,
                    "is_owner": false,
                    "is_system_data": false,
                    "protection": "None",
                });
                json_to_py(py, &default_actions)
            }
        }
    }

    // ========================================================================
    // Lifecycle
    // ========================================================================

    /// Sync all stores to disk.
    fn sync_all(&self) -> PyResult<()> {
        self.inner
            .sync_all()
            .map_err(federation_error_to_pyerr)
    }

    /// Close all stores.
    fn close(&self) -> PyResult<()> {
        self.inner.close().map_err(federation_error_to_pyerr)
    }

    /// Rebuild cross-database index.
    fn rebuild_index(&self) -> PyResult<()> {
        self.inner
            .rebuild_index()
            .map_err(federation_error_to_pyerr)
    }
}

// ============================================================================
// T27.6: FederatedTransaction class
// ============================================================================

/// FederatedTransaction — Cross-store transaction with 2PC.
///
/// Usage:
///     tx = federation.begin_transaction(["social", "docs"])
///     tx.write("social", {"op": "CreateEntity", ...})
///     tx.write("docs", {"op": "CreateDocument", ...})
///     tx.commit()
#[pyclass(name = "FederatedTransaction")]
pub struct PyFederatedTransaction {
    inner: Option<aifastdb_core::federation::FederatedTransaction>,
}

#[pymethods]
impl PyFederatedTransaction {
    /// Add a write operation to the transaction.
    fn write(&mut self, store_alias: &str, operation: &Bound<'_, PyDict>) -> PyResult<()> {
        let tx = self
            .inner
            .as_mut()
            .ok_or_else(|| PyRuntimeError::new_err("Transaction already consumed (committed or rolled back)"))?;
        let op = parse_write_op(operation)?;
        tx.write(store_alias, op)
            .map_err(federation_error_to_pyerr)
    }

    /// Commit the transaction (2PC protocol).
    fn commit(&mut self) -> PyResult<()> {
        let tx = self
            .inner
            .take()
            .ok_or_else(|| PyRuntimeError::new_err("Transaction already consumed"))?;
        tx.commit().map_err(federation_error_to_pyerr)
    }

    /// Rollback the transaction.
    fn rollback(&mut self) -> PyResult<()> {
        let tx = self
            .inner
            .take()
            .ok_or_else(|| PyRuntimeError::new_err("Transaction already consumed"))?;
        tx.rollback().map_err(federation_error_to_pyerr)
    }
}

// ============================================================================
// Helpers
// ============================================================================

/// Convert a MemoryNode to a Python dict.
fn node_to_py(py: Python<'_>, node: &MemoryNode) -> PyResult<PyObject> {
    let val = serde_json::json!({
        "id": node.id,
        "content": node.content,
        "content_type": format!("{:?}", node.content_type),
        "tags": node.tags,
        "metadata": node.metadata,
        "importance": node.importance,
        "created_at": node.created_at,
        "last_accessed_at": node.last_accessed_at,
        "access_count": node.access_count,
        "parent_id": node.parent_id,
    });
    json_to_py(py, &val)
}

/// Convert a slice of MemoryNode to a Vec<PyObject>.
fn nodes_to_py_vec(py: Python<'_>, nodes: &[MemoryNode]) -> PyResult<Vec<PyObject>> {
    let mut results = Vec::with_capacity(nodes.len());
    for node in nodes {
        results.push(node_to_py(py, node)?);
    }
    Ok(results)
}

/// Convert a federation Record to a flat Python dict.
///
/// Each Record variant is flattened with a `_type` field indicating the variant.
fn record_to_py(py: Python<'_>, record: &Record) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    match record {
        Record::GraphEntity { id, entity_type, name, properties, created_at, .. } => {
            dict.set_item("_type", "graph_entity")?;
            dict.set_item("id", id)?;
            dict.set_item("entity_type", entity_type)?;
            dict.set_item("name", name)?;
            dict.set_item("created_at", *created_at)?;
            let props = PyDict::new_bound(py);
            for (k, v) in properties {
                props.set_item(k, json_value_to_py(py, v)?)?;
            }
            dict.set_item("properties", props)?;
        }
        Record::GraphRelation { id, source, target, relation_type, metadata, .. } => {
            dict.set_item("_type", "graph_relation")?;
            dict.set_item("id", id)?;
            dict.set_item("source", source)?;
            dict.set_item("target", target)?;
            dict.set_item("relation_type", relation_type)?;
            let meta = PyDict::new_bound(py);
            for (k, v) in metadata {
                meta.set_item(k, json_value_to_py(py, v)?)?;
            }
            dict.set_item("metadata", meta)?;
        }
        Record::Document { id, doc_type, content, tags, created_at, updated_at, .. } => {
            dict.set_item("_type", "document")?;
            dict.set_item("id", id)?;
            dict.set_item("doc_type", doc_type)?;
            dict.set_item("content", json_value_to_py(py, content)?)?;
            let tag_list: Vec<&str> = tags.iter().map(|t| t.as_str()).collect();
            dict.set_item("tags", tag_list)?;
            dict.set_item("created_at", *created_at)?;
            dict.set_item("updated_at", *updated_at)?;
        }
        Record::Vector { id, vector, metadata, score, .. } => {
            dict.set_item("_type", "vector")?;
            dict.set_item("id", id)?;
            dict.set_item("vector", vector.clone())?;
            let meta = PyDict::new_bound(py);
            for (k, v) in metadata {
                meta.set_item(k, json_value_to_py(py, v)?)?;
            }
            dict.set_item("metadata", meta)?;
            dict.set_item("score", *score)?;
        }
    }
    Ok(dict.into())
}

/// Convert a QueryResult to a Python dict with flattened records.
fn query_result_to_py(py: Python<'_>, result: &QueryResult) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    let records: Vec<PyObject> = result
        .records
        .iter()
        .map(|r| record_to_py(py, r))
        .collect::<PyResult<Vec<_>>>()?;
    dict.set_item("records", records)?;
    dict.set_item("total_count", result.total_count)?;
    dict.set_item("has_more", result.has_more)?;
    let stats_val = serde_json::to_value(&result.stats)
        .map_err(|e| PyRuntimeError::new_err(format!("Stats serialization: {}", e)))?;
    dict.set_item("stats", json_to_py(py, &stats_val)?)?;
    Ok(dict.into())
}

/// Parse a Python dict into a UniversalQuery.
///
/// Handles missing optional fields with sensible defaults:
/// - `filters`: defaults to empty list
/// - `order_by`, `limit`, `offset`, `traverse`, `vector_search`: defaults to None
fn parse_universal_query(dict: &Bound<'_, PyDict>) -> PyResult<UniversalQuery> {
    let json = pydict_to_json(dict)?;
    let obj = json
        .as_object()
        .ok_or_else(|| PyValueError::new_err("Query must be a dict"))?;

    let target = obj
        .get("target")
        .and_then(|v| v.as_str())
        .unwrap_or("*")
        .to_string();

    let filters: Vec<Filter> = obj
        .get("filters")
        .and_then(|v| serde_json::from_value(v.clone()).ok())
        .unwrap_or_default();

    let order_by: Option<(String, SortOrder)> = obj
        .get("order_by")
        .and_then(|v| serde_json::from_value(v.clone()).ok());

    let limit: Option<usize> = obj
        .get("limit")
        .and_then(|v| v.as_u64())
        .map(|v| v as usize);

    let offset: Option<usize> = obj
        .get("offset")
        .and_then(|v| v.as_u64())
        .map(|v| v as usize);

    let traverse = obj
        .get("traverse")
        .and_then(|v| serde_json::from_value(v.clone()).ok());

    let vector_search = obj
        .get("vector_search")
        .and_then(|v| serde_json::from_value(v.clone()).ok());

    Ok(UniversalQuery {
        target,
        filters,
        order_by,
        limit,
        offset,
        traverse,
        vector_search,
    })
}

/// Parse a Python dict into a WriteOp.
///
/// The dict must contain an "op" key specifying the operation type.
fn parse_write_op(dict: &Bound<'_, PyDict>) -> PyResult<WriteOp> {
    let json = pydict_to_json(dict)?;
    let obj = json
        .as_object()
        .ok_or_else(|| PyValueError::new_err("Write operation must be a dict"))?;

    let op_type = obj
        .get("op")
        .and_then(|v| v.as_str())
        .ok_or_else(|| PyValueError::new_err("Write operation must have an 'op' key (e.g. 'CreateDocument')"))?;

    match op_type {
        "CreateDocument" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateDocument requires 'id'"))?
                .to_string();
            let doc_type = obj
                .get("doc_type")
                .and_then(|v| v.as_str())
                .unwrap_or("text")
                .to_string();
            let content = obj
                .get("content")
                .cloned()
                .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
            let tags: Vec<String> = obj
                .get("tags")
                .and_then(|v| serde_json::from_value(v.clone()).ok())
                .unwrap_or_default();
            let owner_id = obj
                .get("owner_id")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            Ok(WriteOp::CreateDocument {
                id,
                doc_type,
                content,
                tags,
                owner_id,
                visibility: None,
            })
        }
        "UpdateDocument" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("UpdateDocument requires 'id'"))?
                .to_string();
            let content = obj.get("content").cloned();
            let tags: Option<Vec<String>> = obj
                .get("tags")
                .and_then(|v| serde_json::from_value(v.clone()).ok());

            Ok(WriteOp::UpdateDocument {
                id,
                content,
                tags,
                caller_id: None,
            })
        }
        "DeleteDocument" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("DeleteDocument requires 'id'"))?
                .to_string();

            Ok(WriteOp::DeleteDocument {
                id,
                caller_id: None,
            })
        }
        "CreateEntity" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateEntity requires 'id'"))?
                .to_string();
            let entity_type = obj
                .get("entity_type")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateEntity requires 'entity_type'"))?
                .to_string();
            let name = obj
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();
            let attributes: HashMap<String, serde_json::Value> = obj
                .get("attributes")
                .and_then(|v| serde_json::from_value(v.clone()).ok())
                .unwrap_or_default();
            let owner_id = obj
                .get("owner_id")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            Ok(WriteOp::CreateEntity {
                entity_type,
                id,
                name,
                attributes,
                owner_id,
                visibility: None,
            })
        }
        "UpdateEntity" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("UpdateEntity requires 'id'"))?
                .to_string();
            let name = obj
                .get("name")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());
            let attributes: Option<HashMap<String, serde_json::Value>> = obj
                .get("attributes")
                .and_then(|v| serde_json::from_value(v.clone()).ok());

            Ok(WriteOp::UpdateEntity {
                id,
                name,
                attributes,
                caller_id: None,
            })
        }
        "DeleteEntity" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("DeleteEntity requires 'id'"))?
                .to_string();

            Ok(WriteOp::DeleteEntity {
                id,
                caller_id: None,
            })
        }
        "CreateRelation" => {
            let source_id = obj
                .get("source_id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateRelation requires 'source_id'"))?
                .to_string();
            let target_id = obj
                .get("target_id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateRelation requires 'target_id'"))?
                .to_string();
            let relation_type = obj
                .get("relation_type")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("CreateRelation requires 'relation_type'"))?
                .to_string();
            let attributes: HashMap<String, serde_json::Value> = obj
                .get("attributes")
                .and_then(|v| serde_json::from_value(v.clone()).ok())
                .unwrap_or_default();
            let owner_id = obj
                .get("owner_id")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string());

            Ok(WriteOp::CreateRelation {
                source_id,
                target_id,
                relation_type,
                attributes,
                owner_id,
            })
        }
        "DeleteRelation" => {
            let relation_id = obj
                .get("relation_id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("DeleteRelation requires 'relation_id'"))?
                .to_string();

            Ok(WriteOp::DeleteRelation {
                relation_id,
                caller_id: None,
            })
        }
        // ===== Vector Operations =====
        "UpsertVector" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("UpsertVector requires 'id'"))?
                .to_string();

            let vector: Vec<f32> = obj
                .get("vector")
                .and_then(|v| v.as_array())
                .ok_or_else(|| PyValueError::new_err("UpsertVector requires 'vector' (array of floats)"))?
                .iter()
                .map(|v| v.as_f64().unwrap_or(0.0) as f32)
                .collect();

            let metadata: HashMap<String, serde_json::Value> = obj
                .get("metadata")
                .and_then(|v| v.as_object())
                .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
                .unwrap_or_default();

            Ok(WriteOp::UpsertVector {
                id,
                vector,
                metadata,
                owner_id: None,
                visibility: None,
            })
        }
        "DeleteVector" => {
            let id = obj
                .get("id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| PyValueError::new_err("DeleteVector requires 'id'"))?
                .to_string();

            Ok(WriteOp::DeleteVector {
                id,
                caller_id: None,
            })
        }
        other => Err(PyValueError::new_err(format!(
            "Unknown write operation type: '{}'. Supported: CreateDocument, UpdateDocument, \
             DeleteDocument, CreateEntity, UpdateEntity, DeleteEntity, CreateRelation, \
             DeleteRelation, UpsertVector, DeleteVector",
            other
        ))),
    }
}

// ============================================================================
// Module Registration
// ============================================================================

impl PyFederation {
    /// Lazily create a tokio runtime for async store adapters.
    ///
    /// After creating the runtime we immediately call `rt.enter()` and
    /// **leak** the `EnterGuard` so that `Handle::current()` is available
    /// for ALL subsequent calls on this Python thread (write, query, close …).
    /// The guard is cleaned up automatically when the `Runtime` itself is dropped.
    fn ensure_runtime(&mut self) -> PyResult<()> {
        if self.runtime.is_none() {
            let rt = tokio::runtime::Builder::new_multi_thread()
                .worker_threads(2)
                .enable_all()
                .build()
                .map_err(|e| PyRuntimeError::new_err(
                    format!("Failed to create async runtime: {}", e),
                ))?;
            // Enter the runtime context permanently for this Federation instance.
            // This makes Handle::current() return this runtime's handle on the
            // current thread, which the core adapters' block_on_async() relies on.
            let guard = rt.enter();
            std::mem::forget(guard);
            self.runtime = Some(rt);
        }
        Ok(())
    }
}

/// Register Federation and DocumentStore classes into the Python module.
pub fn register_federation(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyDocumentStore>()?;
    m.add_class::<PyFederation>()?;
    m.add_class::<PyFederatedTransaction>()?;
    Ok(())
}
