//! Python bindings for high-concurrency storage engines.
//!
//! - `ConcurrentDocumentStore` wraps `aifastdb_core::storage::ConcurrentDocumentStore`
//! - `VectorIndex` wraps `aifastdb_core::vector::ConcurrentHNSWIndex`
//!
//! Both use a dedicated tokio Runtime for async→sync bridging.

use std::sync::Arc;

use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::exceptions::{PyValueError, PyRuntimeError, PyOSError};
use numpy::{PyArray1, IntoPyArray};

use aifastdb_core::storage::concurrent_store::{
    ConcurrentDocumentStore, ConcurrentStoreConfig, ConcurrentStoreError,
};
use aifastdb_core::storage::concurrent::MetricsSnapshot;
use aifastdb_core::types::memory::{MemoryNode, ContentType};
use aifastdb_core::vector::concurrent::{
    ConcurrentHNSWIndex, ConcurrentHNSWConfig,
    VectorSearchResult, VectorShardError,
};
use aifastdb_core::types::config::HNSWConfig;

use crate::types::json_value_to_py;

// ============================================================================
// Helper: tokio Runtime for async bridging
// ============================================================================

fn new_runtime() -> PyResult<tokio::runtime::Runtime> {
    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)))
}

fn concurrent_err(e: ConcurrentStoreError) -> PyErr {
    PyRuntimeError::new_err(format!("ConcurrentStore error: {}", e))
}

fn vector_err(e: VectorShardError) -> PyErr {
    PyRuntimeError::new_err(format!("VectorIndex error: {}", e))
}

// ============================================================================
// Helper: MemoryNode → Python dict (reused from database.rs pattern)
// ============================================================================

fn node_to_dict(py: Python<'_>, node: &MemoryNode) -> PyResult<PyObject> {
    let dict = PyDict::new_bound(py);
    dict.set_item("id", &node.id)?;
    dict.set_item("content", &node.content)?;
    dict.set_item("content_type", format!("{:?}", node.content_type))?;
    dict.set_item("created_at", node.created_at)?;
    dict.set_item("last_accessed_at", node.last_accessed_at)?;
    dict.set_item("access_count", node.access_count)?;
    dict.set_item("importance", node.importance)?;

    if !node.tags.is_empty() {
        dict.set_item("tags", &node.tags)?;
    }
    if let Some(ref parent_id) = node.parent_id {
        dict.set_item("parent_id", parent_id)?;
    }
    if !node.metadata.is_empty() {
        let meta = PyDict::new_bound(py);
        for (k, v) in &node.metadata {
            meta.set_item(k, json_value_to_py(py, v)?)?;
        }
        dict.set_item("metadata", meta)?;
    }

    Ok(dict.into_py(py))
}

fn parse_content_type(s: &str) -> ContentType {
    match s.to_lowercase().as_str() {
        "text" => ContentType::Text,
        "image" => ContentType::Image,
        "code" => ContentType::Code,
        "conversation" => ContentType::Conversation,
        "document" => ContentType::Document,
        _ => ContentType::Other,
    }
}

// ============================================================================
// ConcurrentDocumentStore
// ============================================================================

/// High-concurrency document store with sharded write queues, WAL, and backpressure.
///
/// This is the high-throughput variant of `DocumentStore`. It provides async
/// write buffering, batch aggregation, and write-ahead logging for durability.
///
/// Args:
///     path: Directory path for the store data.
///     shard_count: Number of write shards (0 = auto based on CPU count).
///     enable_wal: Whether to enable Write-Ahead Logging.
///     mode: Performance preset — "default", "high_throughput", or "low_latency".
#[pyclass(name = "ConcurrentDocumentStore")]
pub struct PyConcurrentDocumentStore {
    inner: Arc<ConcurrentDocumentStore>,
    runtime: tokio::runtime::Runtime,
}

#[pymethods]
impl PyConcurrentDocumentStore {
    #[new]
    #[pyo3(signature = (path, *, shard_count=0, enable_wal=true, mode="default"))]
    fn new(
        path: &str,
        shard_count: usize,
        enable_wal: bool,
        mode: &str,
    ) -> PyResult<Self> {
        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 runtime = new_runtime()?;
        let store = runtime.block_on(ConcurrentDocumentStore::new(config))
            .map_err(|e| PyOSError::new_err(format!("Failed to create store: {}", e)))?;

        Ok(Self {
            inner: Arc::new(store),
            runtime,
        })
    }

    // ---- Write API ----

    /// Write a document asynchronously (queued, returns immediately).
    #[pyo3(signature = (id, content, *, tags=vec![], doc_type="text", importance=0.5))]
    fn write(
        &self,
        id: &str,
        content: &str,
        tags: Vec<String>,
        doc_type: &str,
        importance: f64,
    ) -> PyResult<()> {
        let node = Self::make_node(id, content, &tags, doc_type, importance);
        self.runtime.block_on(self.inner.write(node))
            .map(|_| ())
            .map_err(concurrent_err)
    }

    /// Write a document synchronously (waits for persistence).
    #[pyo3(signature = (id, content, *, tags=vec![], doc_type="text", importance=0.5))]
    fn write_sync(
        &self,
        id: &str,
        content: &str,
        tags: Vec<String>,
        doc_type: &str,
        importance: f64,
    ) -> PyResult<()> {
        let node = Self::make_node(id, content, &tags, doc_type, importance);
        self.runtime.block_on(self.inner.write_sync(node))
            .map(|_| ())
            .map_err(concurrent_err)
    }

    /// Write multiple documents in batch.
    ///
    /// Args:
    ///     docs: List of dicts with keys: id, content, tags (opt), doc_type (opt), importance (opt).
    fn write_batch(&self, docs: Vec<Bound<'_, PyDict>>) -> PyResult<usize> {
        let mut nodes = Vec::with_capacity(docs.len());
        for d in &docs {
            let id: String = d.get_item("id")?
                .ok_or_else(|| PyValueError::new_err("Each doc must have 'id'"))?
                .extract()?;
            let content: String = d.get_item("content")?
                .ok_or_else(|| PyValueError::new_err("Each doc must have 'content'"))?
                .extract()?;
            let tags: Vec<String> = d.get_item("tags")?
                .map(|v| v.extract().unwrap_or_default())
                .unwrap_or_default();
            let doc_type: String = d.get_item("doc_type")?
                .map(|v| v.extract().unwrap_or("text".to_string()))
                .unwrap_or("text".to_string());
            let importance: f64 = d.get_item("importance")?
                .map(|v| v.extract().unwrap_or(0.5))
                .unwrap_or(0.5);
            nodes.push(Self::make_node(&id, &content, &tags, &doc_type, importance));
        }
        let count = nodes.len();
        self.runtime.block_on(self.inner.write_batch(nodes))
            .map(|_| count)
            .map_err(concurrent_err)
    }

    // ---- Read API ----

    /// Get a document by ID.
    fn get(&self, py: Python<'_>, id: &str) -> PyResult<Option<PyObject>> {
        match self.inner.get(id) {
            Some(node) => Ok(Some(node_to_dict(py, &node)?)),
            None => Ok(None),
        }
    }

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

    fn __len__(&self) -> usize {
        self.inner.len()
    }

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

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

    // ---- Management ----

    /// Flush all pending writes to disk.
    fn flush(&self) -> PyResult<()> {
        self.runtime.block_on(self.inner.flush())
            .map_err(concurrent_err)
    }

    /// Get performance metrics.
    fn metrics(&self, py: Python<'_>) -> PyResult<PyObject> {
        let m = self.inner.metrics();
        let dict = PyDict::new_bound(py);
        dict.set_item("total_requests", m.total_requests)?;
        dict.set_item("successful_writes", m.successful_writes)?;
        dict.set_item("failed_writes", m.failed_writes)?;
        dict.set_item("rejected_requests", m.rejected_requests)?;
        dict.set_item("batches_flushed", m.batches_flushed)?;
        dict.set_item("batch_documents", m.batch_documents)?;
        dict.set_item("current_queue_size", m.current_queue_size)?;
        dict.set_item("peak_queue_size", m.peak_queue_size)?;
        dict.set_item("bytes_written", m.bytes_written)?;
        dict.set_item("wal_bytes_written", m.wal_bytes_written)?;
        dict.set_item("avg_queue_time_us", m.avg_queue_time_us)?;
        dict.set_item("avg_write_time_us", m.avg_write_time_us)?;
        dict.set_item("success_rate", m.success_rate())?;
        Ok(dict.into_py(py))
    }

    /// Get current queue fill ratio (0.0 ~ 1.0).
    fn queue_fill_ratio(&self) -> f64 {
        self.inner.queue_fill_ratio()
    }

    /// Recover documents from WAL after a crash.
    fn recover(&self, py: Python<'_>) -> PyResult<PyObject> {
        let r = self.inner.recover().map_err(concurrent_err)?;
        let dict = PyDict::new_bound(py);
        dict.set_item("entries_found", r.entries_found)?;
        dict.set_item("entries_replayed", r.entries_replayed)?;
        dict.set_item("entries_skipped", r.entries_skipped)?;
        dict.set_item("entries_failed", r.entries_failed)?;
        dict.set_item("last_sequence", r.last_sequence)?;
        Ok(dict.into_py(py))
    }

    /// Number of write shards.
    #[getter]
    fn shard_count(&self) -> usize {
        self.inner.shard_count()
    }

    /// Gracefully shutdown the store.
    fn close(&self) -> PyResult<()> {
        self.runtime.block_on(self.inner.shutdown())
            .map_err(concurrent_err)
    }

    fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    #[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.runtime.block_on(self.inner.flush());
        let _ = self.runtime.block_on(self.inner.shutdown());
        Ok(false)
    }

    fn __repr__(&self) -> String {
        format!(
            "ConcurrentDocumentStore(docs={}, shards={})",
            self.inner.len(),
            self.inner.shard_count(),
        )
    }
}

impl PyConcurrentDocumentStore {
    fn make_node(id: &str, content: &str, tags: &[String], doc_type: &str, importance: f64) -> MemoryNode {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as i64;
        MemoryNode {
            id: id.to_string(),
            content: content.to_string(),
            embedding: None,
            content_type: parse_content_type(doc_type),
            parent_id: None,
            children: vec![],
            depth: 0,
            metadata: Default::default(),
            tags: tags.to_vec(),
            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,
        }
    }
}

// ============================================================================
// VectorIndex (ConcurrentHNSWIndex)
// ============================================================================

/// High-concurrency sharded HNSW vector index.
///
/// Provides async write queuing, batch aggregation, cross-shard parallel search,
/// and backpressure control for high-throughput vector operations.
///
/// Args:
///     dimension: Vector dimension.
///     shard_count: Number of shards (0 = auto based on CPU count).
///     mode: Performance preset — "default", "high_throughput", or "low_latency".
///     m: HNSW M parameter (default 16).
///     ef_construction: HNSW ef_construction (default 100).
///     ef_search: HNSW ef_search (default 50).
///     max_elements: Maximum number of vectors (default 100000).
#[pyclass(name = "VectorIndex")]
pub struct PyVectorIndex {
    inner: ConcurrentHNSWIndex,
    runtime: tokio::runtime::Runtime,
    dim: usize,
}

#[pymethods]
impl PyVectorIndex {
    #[new]
    #[pyo3(signature = (dimension, *, shard_count=0, mode="default", m=16, ef_construction=100, ef_search=50, max_elements=100000))]
    fn new(
        dimension: usize,
        shard_count: usize,
        mode: &str,
        m: usize,
        ef_construction: usize,
        ef_search: usize,
        max_elements: usize,
    ) -> PyResult<Self> {
        if dimension == 0 {
            return Err(PyValueError::new_err("dimension must be > 0"));
        }

        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 runtime = new_runtime()?;
        let index = ConcurrentHNSWIndex::new(config);

        Ok(Self {
            inner: index,
            runtime,
            dim: dimension,
        })
    }

    // ---- Write API ----

    /// Insert a vector (queued, returns immediately).
    fn insert(&self, id: &str, vector: Vec<f32>) -> PyResult<()> {
        self.check_dim(vector.len())?;
        self.runtime.block_on(self.inner.insert(id.to_string(), vector))
            .map(|_| ())
            .map_err(vector_err)
    }

    /// Insert a vector and wait for completion.
    fn insert_sync(&self, id: &str, vector: Vec<f32>) -> PyResult<()> {
        self.check_dim(vector.len())?;
        self.runtime.block_on(self.inner.insert_sync(id.to_string(), vector))
            .map(|_| ())
            .map_err(vector_err)
    }

    /// Batch insert multiple vectors.
    ///
    /// Args:
    ///     items: List of (id, vector) tuples.
    fn insert_batch(&self, items: Vec<(String, Vec<f32>)>) -> PyResult<usize> {
        for (_, v) in &items {
            self.check_dim(v.len())?;
        }
        let count = items.len();
        self.runtime.block_on(self.inner.insert_batch(items))
            .map(|_| count)
            .map_err(vector_err)
    }

    /// Delete a vector by ID.
    fn delete(&self, id: &str) -> PyResult<()> {
        self.runtime.block_on(self.inner.delete(id))
            .map(|_| ())
            .map_err(vector_err)
    }

    // ---- Search API ----

    /// Search for k nearest neighbors.
    ///
    /// Returns list of dicts: [{"id": str, "score": float}, ...]
    #[pyo3(signature = (query, k=10))]
    fn search(&self, py: Python<'_>, query: Vec<f32>, k: usize) -> PyResult<Vec<PyObject>> {
        self.check_dim(query.len())?;
        let results = self.inner.search(&query, k);
        self.results_to_py(py, results)
    }

    /// Search with custom ef parameter.
    #[pyo3(signature = (query, k=10, ef=200))]
    fn search_with_ef(&self, py: Python<'_>, query: Vec<f32>, k: usize, ef: usize) -> PyResult<Vec<PyObject>> {
        self.check_dim(query.len())?;
        let results = self.inner.search_with_ef(&query, k, ef);
        self.results_to_py(py, results)
    }

    /// Search using a numpy array as query.
    #[pyo3(signature = (query, k=10))]
    fn search_numpy(&self, py: Python<'_>, query: numpy::PyReadonlyArray1<f32>, k: usize) -> PyResult<Vec<PyObject>> {
        let vec = query.as_slice()?.to_vec();
        self.check_dim(vec.len())?;
        let results = self.inner.search(&vec, k);
        self.results_to_py(py, results)
    }

    /// Get a vector by ID. Returns list[float] or None.
    fn get(&self, id: &str) -> Option<Vec<f32>> {
        self.inner.get(id).map(|v| v.as_ref().clone())
    }

    /// Get a vector as numpy array. Returns ndarray or None.
    fn get_numpy<'py>(&self, py: Python<'py>, id: &str) -> Option<Bound<'py, PyArray1<f32>>> {
        self.inner.get(id).map(|v| {
            v.as_ref().clone().into_pyarray_bound(py)
        })
    }

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

    fn __len__(&self) -> usize {
        self.inner.len()
    }

    fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    // ---- Management ----

    /// Flush all pending writes.
    fn flush(&self) -> PyResult<()> {
        self.runtime.block_on(self.inner.flush())
            .map_err(vector_err)
    }

    /// Get performance metrics.
    fn metrics(&self, py: Python<'_>) -> PyResult<PyObject> {
        let m = self.inner.metrics();
        let dict = PyDict::new_bound(py);
        dict.set_item("total_requests", m.total_requests)?;
        dict.set_item("successful_writes", m.successful_writes)?;
        dict.set_item("failed_writes", m.failed_writes)?;
        dict.set_item("rejected_requests", m.rejected_requests)?;
        dict.set_item("batches_flushed", m.batches_flushed)?;
        dict.set_item("batch_documents", m.batch_documents)?;
        dict.set_item("current_queue_size", m.current_queue_size)?;
        dict.set_item("peak_queue_size", m.peak_queue_size)?;
        dict.set_item("bytes_written", m.bytes_written)?;
        dict.set_item("avg_queue_time_us", m.avg_queue_time_us)?;
        dict.set_item("avg_write_time_us", m.avg_write_time_us)?;
        dict.set_item("success_rate", m.success_rate())?;
        Ok(dict.into_py(py))
    }

    /// Get queue fill ratio (0.0 ~ 1.0).
    fn queue_fill_ratio(&self) -> f64 {
        self.inner.queue_fill_ratio()
    }

    /// Number of shards.
    #[getter]
    fn shard_count(&self) -> usize {
        self.inner.shard_count()
    }

    /// Per-shard statistics: list of (shard_id, vector_count) tuples.
    fn shard_stats(&self) -> Vec<(usize, usize)> {
        self.inner.shard_stats()
    }

    /// Vector dimension.
    #[getter]
    fn dimension(&self) -> usize {
        self.dim
    }

    /// Backpressure status string.
    fn backpressure_status(&self) -> String {
        format!("{:?}", self.inner.backpressure_status())
    }

    fn close(&self) -> PyResult<()> {
        self.runtime.block_on(self.inner.flush())
            .map_err(vector_err)
    }

    fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    #[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.runtime.block_on(self.inner.flush());
        Ok(false)
    }

    fn __repr__(&self) -> String {
        format!(
            "VectorIndex(dim={}, vectors={}, shards={})",
            self.dim,
            self.inner.len(),
            self.inner.shard_count(),
        )
    }
}

impl PyVectorIndex {
    fn check_dim(&self, len: usize) -> PyResult<()> {
        if len != self.dim {
            return Err(PyValueError::new_err(format!(
                "Vector dimension mismatch: expected {}, got {}",
                self.dim, len
            )));
        }
        Ok(())
    }

    fn results_to_py(&self, py: Python<'_>, results: Vec<VectorSearchResult>) -> PyResult<Vec<PyObject>> {
        results
            .into_iter()
            .map(|r| {
                let dict = PyDict::new_bound(py);
                dict.set_item("id", &r.id)?;
                dict.set_item("distance", r.distance)?;
                dict.set_item("shard_id", r.shard_id)?;
                Ok(dict.into_py(py))
            })
            .collect()
    }
}

// ============================================================================
// Module registration
// ============================================================================

pub fn register_concurrent(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyConcurrentDocumentStore>()?;
    m.add_class::<PyVectorIndex>()?;
    Ok(())
}
