//! Python bindings for the Luci search engine via PyO3.

#[cfg(feature = "mimalloc")]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

use luci_index::index::Index as RustIndex;
use luci_mapping::Mapping;

/// A Luci search index.
///
/// Create or open an index, add documents, commit, and search.
///
/// Example::
///
///     index = Index.create("data.luci", {
///         "properties": {
///             "title": {"type": "text"},
///             "tag": {"type": "keyword"},
///         }
///     })
///     index.add({"title": "hello world", "tag": "greeting"})
///     index.commit()
///     results = index.search({"match": {"title": "hello"}})
///
///     # Or fully dynamic (no mapping needed):
///     index = Index.create("data.luci")
///
///     # Open existing (mappings loaded from disk):
///     index = Index.open("data.luci")
#[pyclass(unsendable)]
struct Index {
    inner: RustIndex,
}

#[pymethods]
impl Index {
    /// Create a new index at the given path.
    ///
    /// Args:
    ///     path: File path for the index (e.g. ``"data.luci"``).
    ///     mapping: Optional ES-compatible mapping dict, e.g.
    ///         ``{"properties": {"title": {"type": "text"}}}``.
    ///         If omitted, all field types are inferred dynamically.
    ///     memory_budget: Memory budget in bytes before auto-flush to a new
    ///         segment. Smaller values produce more segments (better query
    ///         parallelism, more merge work). Default: 64 MB (67108864).
    ///         Not persisted — applies only to this session's writes.
    #[staticmethod]
    #[pyo3(signature = (path, mapping=None, memory_budget=None))]
    fn create(path: &str, mapping: Option<PyObject>, memory_budget: Option<usize>, py: Python<'_>) -> PyResult<Self> {
        let mut inner = match mapping {
            Some(m) => {
                let mapping_json = py_to_json(m, py)?;
                let mapping = Mapping::from_json(&mapping_json)
                    .map_err(|e| PyValueError::new_err(format!("invalid mapping: {e}")))?;
                RustIndex::create_with_mapping(path, mapping)
            }
            None => RustIndex::create(path),
        }
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        if let Some(budget) = memory_budget {
            inner.set_memory_budget(budget);
        }
        Ok(Self { inner })
    }

    /// Open an existing index.
    ///
    /// Mappings are loaded from the index file automatically.
    ///
    /// Args:
    ///     path: File path for the index.
    ///     memory_budget: Memory budget in bytes before auto-flush to a new
    ///         segment. Default: 64 MB (67108864).
    ///         Not persisted — applies only to this session's writes.
    #[staticmethod]
    #[pyo3(signature = (path, memory_budget=None))]
    fn open(path: &str, memory_budget: Option<usize>) -> PyResult<Self> {
        let mut inner = RustIndex::open(path)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        if let Some(budget) = memory_budget {
            inner.set_memory_budget(budget);
        }
        Ok(Self { inner })
    }

    /// Add a document to the index.
    ///
    /// The document is buffered until :meth:`commit` is called.
    ///
    /// Args:
    ///     doc: A dict representing the document.
    fn add(&mut self, doc: PyObject, py: Python<'_>) -> PyResult<()> {
        let doc_json = py_to_json(doc, py)?;
        self.inner.add(doc_json)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        Ok(())
    }

    /// Commit buffered documents, making them searchable.
    fn commit(&mut self) -> PyResult<()> {
        self.inner.commit()
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        Ok(())
    }

    /// Force-merge all segments down to at most ``max_segments``.
    ///
    /// Call after bulk indexing to optimize for query performance.
    ///
    /// Args:
    ///     max_segments: Target number of segments (default: 1).
    #[pyo3(signature = (max_segments=1))]
    fn force_merge(&mut self, max_segments: usize) -> PyResult<()> {
        self.inner.force_merge(max_segments)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        Ok(())
    }

    /// Search the index.
    ///
    /// Args:
    ///     query: ES-compatible query dict, e.g.
    ///         ``{"match": {"title": "hello"}}`` or
    ///         ``{"query": {"match": {"title": "hello"}}, "size": 10}``
    ///     size: Maximum number of results (default 10).
    ///
    /// Returns:
    ///     A dict with ``hits``, ``total_hits``, and optionally ``aggregations``.
    #[pyo3(signature = (query, size=10))]
    fn search(&mut self, query: PyObject, size: usize, py: Python<'_>) -> PyResult<PyObject> {
        let query_json = py_to_json(query, py)?;
        let results = self.inner.search(query_json, size)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        // Convert to Python dict
        let hits: Vec<serde_json::Value> = results.hits.iter().map(|h| {
            let mut hit = serde_json::Map::new();
            hit.insert("_score".into(), serde_json::json!(h.score));
            hit.insert("_doc_id".into(), serde_json::json!(h.doc_id.as_u32()));
            if let Some(ref source) = h.source {
                hit.insert("_source".into(), source.clone());
            }
            if let Some(ref fields) = h.fields {
                if !fields.is_empty() {
                    hit.insert("fields".into(), serde_json::Value::Object(fields.clone()));
                }
            }
            if let Some(ref sort_vals) = h.sort_values {
                let arr: Vec<serde_json::Value> = sort_vals.iter().map(|sv| sv.to_json()).collect();
                hit.insert("sort".into(), serde_json::Value::Array(arr));
            }
            if let Some(ref explanation) = h.explanation {
                hit.insert("_explanation".into(), explanation.to_json());
            }
            if let Some(ref inner) = h.inner_hits {
                let ih: serde_json::Map<String, serde_json::Value> = inner.iter()
                    .map(|(k, v)| (k.clone(), v.clone()))
                    .collect();
                hit.insert("inner_hits".into(), serde_json::Value::Object(ih));
            }
            serde_json::Value::Object(hit)
        }).collect();

        let mut result = serde_json::Map::new();
        result.insert("hits".into(), serde_json::Value::Array(hits));
        result.insert("total_hits".into(), results.total_hits.to_json());

        if !results.aggregations.is_empty() {
            let mut aggs = serde_json::Map::new();
            for (name, agg_result) in &results.aggregations {
                aggs.insert(name.clone(), agg_result.to_json());
            }
            result.insert("aggregations".into(), serde_json::Value::Object(aggs));
        }

        json_to_py(&serde_json::Value::Object(result), py)
    }

    /// Get a document by its `_id`. Returns None if not found.
    fn get(&mut self, id: &str, py: Python<'_>) -> PyResult<Option<PyObject>> {
        let hit = self.inner.get(id)
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;
        match hit {
            Some(h) => {
                let mut obj = serde_json::Map::new();
                obj.insert("_id".into(), serde_json::json!(id));
                if let Some(source) = h.source {
                    obj.insert("_source".into(), source);
                }
                json_to_py(&serde_json::Value::Object(obj), py).map(Some)
            }
            None => Ok(None),
        }
    }

    /// Delete a document by its `_id`. Returns true if found and deleted.
    fn delete(&mut self, id: &str) -> PyResult<bool> {
        self.inner.delete(id)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// Update a document by its `_id` with a partial doc merge.
    fn update(&mut self, id: &str, doc: PyObject, py: Python<'_>) -> PyResult<bool> {
        let partial = py_to_json(doc, py)?;
        self.inner.update(id, partial)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// Delete all documents matching a query. Returns count deleted.
    fn delete_by_query(&mut self, query: PyObject, py: Python<'_>) -> PyResult<u64> {
        let query_json = py_to_json(query, py)?;
        self.inner.delete_by_query(query_json)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// Count documents matching a query.
    fn count(&mut self, query: PyObject, py: Python<'_>) -> PyResult<u64> {
        let query_json = py_to_json(query, py)?;
        self.inner.count(query_json)
            .map_err(|e| PyValueError::new_err(format!("{e}")))
    }

    /// Number of buffered (uncommitted) documents.
    fn buffered_doc_count(&self) -> u32 {
        self.inner.buffered_doc_count()
    }
}

/// Convert a Python object to serde_json::Value via JSON serialization.
fn py_to_json(obj: PyObject, py: Python<'_>) -> PyResult<serde_json::Value> {
    let json_mod = py.import("json")?;
    let json_str: String = json_mod.call_method1("dumps", (obj,))?.extract()?;
    serde_json::from_str(&json_str)
        .map_err(|e| PyValueError::new_err(format!("invalid JSON: {e}")))
}

/// Convert a serde_json::Value to a Python object via JSON serialization.
fn json_to_py(val: &serde_json::Value, py: Python<'_>) -> PyResult<PyObject> {
    let json_str = serde_json::to_string(val)
        .map_err(|e| PyValueError::new_err(format!("JSON serialization failed: {e}")))?;
    let json_mod = py.import("json")?;
    let result = json_mod.call_method1("loads", (json_str,))?;
    Ok(result.into())
}

/// Configure the number of threads used for parallel search.
///
/// Must be called before any search operation. Can only be called once
/// per process. If not called, defaults to the number of available CPUs.
#[pyfunction]
#[pyo3(signature = (num_threads))]
fn set_num_threads(num_threads: usize) -> PyResult<()> {
    luci_index::set_num_threads(num_threads)
        .map_err(|e| PyValueError::new_err(format!("{e}")))
}

/// Python module initialization.
#[pymodule]
fn _luci(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Index>()?;
    m.add_function(wrap_pyfunction!(set_num_threads, m)?)?;
    Ok(())
}
