//! Python bindings for VectorStore, InMemoryVectorStore, Embeddings, Retriever

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

use flowgentra_ai::core::rag::{
    Embeddings, InMemoryVectorStore, Retriever, RetrievalConfig, RetrieverStrategy,
    VectorStoreBackend,
    OpenAIEmbeddings, HuggingFaceEmbeddings, OllamaEmbeddings, MistralEmbeddings,
    CachedEmbeddings,
};

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

// ─── PyEmbeddings ───────────────────────────────────────────────────────────

/// Embeddings manager for generating vector embeddings from text.
///
/// Example:
///     emb = Embeddings.mock(128)
///     vector = emb.embed("Hello world")
///     print(len(vector))  # 128
#[pyclass(name = "Embeddings")]
pub struct PyEmbeddings {
    pub(crate) inner: Arc<Embeddings>,
}

#[pymethods]
impl PyEmbeddings {
    /// Create mock embeddings for testing (hash-based, no API needed).
    #[staticmethod]
    fn mock(dimension: usize) -> Self {
        PyEmbeddings {
            inner: Arc::new(Embeddings::mock(dimension)),
        }
    }

    /// Create OpenAI embeddings.
    ///
    /// Example:
    ///     emb = Embeddings.openai("sk-...", "text-embedding-3-small")
    #[staticmethod]
    #[pyo3(signature = (api_key, model="text-embedding-3-small"))]
    fn openai(api_key: &str, model: &str) -> Self {
        let provider = OpenAIEmbeddings::new(api_key, model);
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(provider))),
        }
    }

    /// Create OpenAI embeddings with custom dimension.
    #[staticmethod]
    fn openai_with_dimension(api_key: &str, model: &str, dimension: usize) -> Self {
        let provider = OpenAIEmbeddings::new(api_key, model).with_dimension(dimension);
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(provider))),
        }
    }

    /// Create HuggingFace embeddings.
    #[staticmethod]
    #[pyo3(signature = (model, api_key, endpoint=None, dimension=None))]
    fn huggingface(model: &str, api_key: &str, endpoint: Option<&str>, dimension: Option<usize>) -> Self {
        let mut provider = HuggingFaceEmbeddings::new(model, api_key);
        if let Some(ep) = endpoint {
            provider = provider.with_endpoint(ep);
        }
        if let Some(dim) = dimension {
            provider = provider.with_dimension(dim);
        }
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(provider))),
        }
    }

    /// Create Ollama embeddings (local).
    #[staticmethod]
    #[pyo3(signature = (model, base_url=None, dimension=None))]
    fn ollama(model: &str, base_url: Option<String>, dimension: Option<usize>) -> Self {
        let mut provider = OllamaEmbeddings::new(model, base_url);
        if let Some(dim) = dimension {
            provider = provider.with_dimension(dim);
        }
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(provider))),
        }
    }

    /// Create Mistral embeddings.
    #[staticmethod]
    #[pyo3(signature = (api_key, model=None))]
    fn mistral(api_key: &str, model: Option<String>) -> Self {
        let provider = MistralEmbeddings::new(api_key, model);
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(provider))),
        }
    }

    /// Create OpenAI embeddings with a cache layer.
    #[staticmethod]
    #[pyo3(signature = (api_key, model="text-embedding-3-small"))]
    fn openai_cached(api_key: &str, model: &str) -> Self {
        let provider = OpenAIEmbeddings::new(api_key, model);
        let cached = CachedEmbeddings::new(Arc::new(provider));
        PyEmbeddings {
            inner: Arc::new(Embeddings::new(Arc::new(cached))),
        }
    }

    /// Generate embedding for a single text.
    fn embed(&self, text: &str) -> PyResult<Vec<f32>> {
        let rt = get_runtime();
        rt.block_on(self.inner.embed(text))
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{}", e)))
    }

    /// Generate embeddings for multiple texts (batch).
    fn embed_batch(&self, texts: Vec<String>) -> PyResult<Vec<Vec<f32>>> {
        let rt = get_runtime();
        let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
        rt.block_on(self.inner.embed_batch(refs))
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{}", e)))
    }

    /// Get embedding dimension.
    fn get_dimension(&self) -> usize {
        self.inner.get_dimension()
    }

    fn __repr__(&self) -> String {
        format!("Embeddings(dim={})", self.inner.get_dimension())
    }
}

// ─── PyInMemoryVectorStore ──────────────────────────────────────────────────

/// In-memory vector store with cosine similarity search.
///
/// Example:
///     store = InMemoryVectorStore()
///     emb = Embeddings.mock(128)
///     doc = Document("doc1", "Hello world")
///     doc_emb = emb.embed("Hello world")
///     store.index(doc, doc_emb)
///     results = store.search(emb.embed("Hello"), top_k=5)
#[pyclass(name = "InMemoryVectorStore")]
pub struct PyInMemoryVectorStore {
    pub(crate) inner: Arc<InMemoryVectorStore>,
}

#[pymethods]
impl PyInMemoryVectorStore {
    #[new]
    fn new() -> Self {
        PyInMemoryVectorStore {
            inner: Arc::new(InMemoryVectorStore::new()),
        }
    }

    /// 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)
    }

    /// 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 {
        "InMemoryVectorStore(...)".to_string()
    }
}

// ─── PyRetrievalConfig ──────────────────────────────────────────────────────

/// Configuration for retrieval strategies.
///
/// Example:
///     config = RetrievalConfig.semantic(top_k=5, threshold=0.7)
///     config = RetrievalConfig.hybrid(keyword_weight=0.3, top_k=10)
#[pyclass(name = "RetrievalConfig")]
#[derive(Clone)]
pub struct PyRetrievalConfig {
    pub(crate) inner: RetrievalConfig,
}

#[pymethods]
impl PyRetrievalConfig {
    /// Create a semantic search config.
    #[staticmethod]
    #[pyo3(signature = (top_k=5, threshold=0.7))]
    fn semantic(top_k: usize, threshold: f32) -> Self {
        PyRetrievalConfig {
            inner: RetrievalConfig::new(RetrieverStrategy::SemanticSearch)
                .with_top_k(top_k)
                .with_threshold(threshold),
        }
    }

    /// Create a hybrid (semantic + keyword) search config.
    #[staticmethod]
    #[pyo3(signature = (keyword_weight=0.3, top_k=5, threshold=0.7))]
    fn hybrid(keyword_weight: f32, top_k: usize, threshold: f32) -> Self {
        PyRetrievalConfig {
            inner: RetrievalConfig::new(RetrieverStrategy::hybrid(keyword_weight))
                .with_top_k(top_k)
                .with_threshold(threshold),
        }
    }

    #[getter]
    fn top_k(&self) -> usize {
        self.inner.top_k
    }

    #[getter]
    fn similarity_threshold(&self) -> f32 {
        self.inner.similarity_threshold
    }

    fn __repr__(&self) -> String {
        format!(
            "RetrievalConfig(top_k={}, threshold={})",
            self.inner.top_k, self.inner.similarity_threshold
        )
    }
}

// ─── PyRetriever ────────────────────────────────────────────────────────────

/// Full retrieval pipeline: embed → search → hybrid → rerank → dedup.
///
/// Example:
///     store = InMemoryVectorStore()
///     emb = Embeddings.mock(128)
///     config = RetrievalConfig.semantic(top_k=3, threshold=0.0)
///     retriever = Retriever(store, emb, config)
///     results = retriever.retrieve("What is Rust?")
#[pyclass(name = "Retriever")]
pub struct PyRetriever {
    inner: Retriever,
}

#[pymethods]
impl PyRetriever {
    #[new]
    fn new(
        store: &PyInMemoryVectorStore,
        embeddings: &PyEmbeddings,
        config: &PyRetrievalConfig,
    ) -> Self {
        PyRetriever {
            inner: Retriever::new(
                store.inner.clone() as Arc<dyn VectorStoreBackend>,
                embeddings.inner.clone(),
                config.inner.clone(),
            ),
        }
    }

    /// Enable deduplication with a similarity threshold (e.g., 0.85).
    fn with_dedup(&mut self, threshold: f32) {
        // Reconstruct since with_dedup consumes self
        let inner = std::mem::replace(
            &mut self.inner,
            Retriever::new(
                Arc::new(InMemoryVectorStore::new()),
                Arc::new(Embeddings::mock(1)),
                RetrievalConfig::default(),
            ),
        );
        self.inner = inner.with_dedup(threshold);
    }

    /// Execute the full retrieval pipeline for a query.
    fn retrieve(&self, query: &str) -> PyResult<Vec<PySearchResult>> {
        let rt = get_runtime();
        let results = rt
            .block_on(self.inner.retrieve(query))
            .map_err(to_py_err_generic)?;
        Ok(results
            .into_iter()
            .map(|r| PySearchResult { inner: r })
            .collect())
    }

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