//! Python bindings for ChromaDB vector store

use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::collections::HashMap;
use std::sync::Arc;

use flowgentra_ai::core::rag::{
    ChromaStore, RAGConfig, VectorStoreBackend, VectorStoreType,
};

use crate::error::to_py_err_generic;
use crate::get_runtime;
use crate::py_to_json;
use crate::rag::{PyDocument, PySearchResult};

// ─── PyChromaStore ──────────────────────────────────────────────────────────

/// ChromaDB vector store backend.
///
/// Connects to a ChromaDB instance via REST API.
///
/// Example:
///     store = ChromaStore("my_collection", endpoint="http://localhost:8000")
///     store.index(doc, embedding)
///     results = store.search(query_embedding, top_k=5)
#[pyclass(name = "ChromaStore")]
pub struct PyChromaStore {
    pub(crate) inner: Arc<ChromaStore>,
}

#[pymethods]
impl PyChromaStore {
    /// Create a new ChromaStore connected to a ChromaDB instance.
    ///
    /// Args:
    ///     collection_name: Name of the collection to use/create.
    ///     endpoint: ChromaDB REST API endpoint (default: "http://localhost:8000").
    ///     embedding_dim: Embedding dimension (default: 1536).
    #[new]
    #[pyo3(signature = (collection_name, endpoint="http://localhost:8000", embedding_dim=1536))]
    fn new(collection_name: &str, endpoint: &str, embedding_dim: usize) -> PyResult<Self> {
        let config = RAGConfig {
            store_type: VectorStoreType::Chroma,
            api_key: None,
            endpoint: Some(endpoint.to_string()),
            index_name: collection_name.to_string(),
            embedding_dim,
        };
        let rt = get_runtime();
        let store = rt
            .block_on(ChromaStore::new(&config))
            .map_err(to_py_err_generic)?;
        Ok(PyChromaStore {
            inner: Arc::new(store),
        })
    }

    /// Index a document with its embedding.
    fn index(&self, doc: &PyDocument, embedding: Vec<f32>) -> PyResult<()> {
        let mut document = doc.inner.clone();
        document.embedding = Some(embedding);
        let rt = get_runtime();
        rt.block_on(self.inner.index(document))
            .map_err(to_py_err_generic)
    }

    /// Search for similar documents by embedding vector.
    #[pyo3(signature = (query_embedding, top_k=5, filter=None))]
    fn search(
        &self,
        query_embedding: Vec<f32>,
        top_k: usize,
        filter: Option<&Bound<'_, PyDict>>,
    ) -> PyResult<Vec<PySearchResult>> {
        let metadata_filter = if let Some(dict) = filter {
            let mut map = HashMap::new();
            for (k, v) in dict.iter() {
                let key: String = k.extract()?;
                let val = py_to_json(&v)?;
                map.insert(key, val);
            }
            Some(map)
        } else {
            None
        };

        let rt = get_runtime();
        let results = rt
            .block_on(self.inner.search(query_embedding, top_k, metadata_filter))
            .map_err(to_py_err_generic)?;

        Ok(results
            .into_iter()
            .map(|r| PySearchResult { inner: r })
            .collect())
    }

    /// Delete a document by ID.
    fn delete(&self, doc_id: &str) -> PyResult<()> {
        let rt = get_runtime();
        rt.block_on(self.inner.delete(doc_id))
            .map_err(to_py_err_generic)
    }

    /// Update a document.
    #[pyo3(signature = (doc, embedding=None))]
    fn update(&self, doc: &PyDocument, embedding: Option<Vec<f32>>) -> PyResult<()> {
        let mut document = doc.inner.clone();
        document.embedding = embedding;
        let rt = get_runtime();
        rt.block_on(self.inner.update(document))
            .map_err(to_py_err_generic)
    }

    /// Get a document by ID.
    fn get(&self, doc_id: &str) -> PyResult<PyDocument> {
        let rt = get_runtime();
        let doc = rt
            .block_on(self.inner.get(doc_id))
            .map_err(to_py_err_generic)?;
        Ok(PyDocument { inner: doc })
    }

    /// List all documents.
    fn list(&self) -> PyResult<Vec<PyDocument>> {
        let rt = get_runtime();
        let docs = rt.block_on(self.inner.list()).map_err(to_py_err_generic)?;
        Ok(docs.into_iter().map(|d| PyDocument { inner: d }).collect())
    }

    /// Clear all documents.
    fn clear(&self) -> PyResult<()> {
        let rt = get_runtime();
        rt.block_on(self.inner.clear()).map_err(to_py_err_generic)
    }

    fn __repr__(&self) -> String {
        "ChromaStore(...)".to_string()
    }
}
