"""
ZeroDB Vector Store for LangChain - Issue #1372

Usage:
    from langchain_zerodb import ZeroDBVectorStore

    vectorstore = ZeroDBVectorStore(
        api_key="your-api-key",
        project_id="your-project-id",
    )

    # Add documents
    vectorstore.add_texts(["Hello world", "ZeroDB is fast"])

    # Similarity search
    results = vectorstore.similarity_search("hello", k=5)
"""

import httpx
from typing import Any, Dict, Iterable, List, Optional, Tuple

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore


class ZeroDBVectorStore(VectorStore):
    """ZeroDB vector store for LangChain.

    Uses ZeroDB's embedding API (TEI-powered, FREE) for vector generation
    and pgvector with HNSW indexes for similarity search.
    """

    def __init__(
        self,
        api_key: str,
        project_id: str,
        base_url: str = "https://api.ainative.studio",
        namespace: str = "default",
        embedding: Optional[Embeddings] = None,
    ):
        self.api_key = api_key
        self.project_id = project_id
        self.base_url = base_url.rstrip("/")
        self.namespace = namespace
        self._embedding = embedding
        self._client = httpx.Client(
            base_url=self.base_url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            timeout=30.0,
        )

    @property
    def embeddings(self) -> Optional[Embeddings]:
        return self._embedding

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Add texts to the vector store using ZeroDB's embed-and-store endpoint."""
        texts_list = list(texts)
        metadatas = metadatas or [{}] * len(texts_list)

        response = self._client.post(
            f"/api/v1/projects/{self.project_id}/embeddings/embed-and-store",
            json={
                "texts": texts_list,
                "namespace": self.namespace,
                "metadata_list": metadatas,
            },
        )
        response.raise_for_status()
        result = response.json()
        return result.get("vector_ids", [])

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        **kwargs: Any,
    ) -> List[Document]:
        """Search for similar documents."""
        results = self.similarity_search_with_score(query, k=k, **kwargs)
        return [doc for doc, _ in results]

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        threshold: float = 0.3,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Search for similar documents with relevance scores."""
        response = self._client.post(
            f"/api/v1/projects/{self.project_id}/embeddings/search",
            json={
                "query": query,
                "limit": k,
                "threshold": threshold,
                "namespace": self.namespace,
            },
        )
        response.raise_for_status()
        data = response.json()

        results = []
        for item in data.get("results", []):
            doc = Document(
                page_content=item.get("document", ""),
                metadata=item.get("metadata", {}),
            )
            score = item.get("similarity", 0.0)
            results.append((doc, score))
        return results

    @classmethod
    def from_texts(
        cls,
        texts: List[str],
        embedding: Optional[Embeddings] = None,
        metadatas: Optional[List[dict]] = None,
        api_key: str = "",
        project_id: str = "",
        **kwargs: Any,
    ) -> "ZeroDBVectorStore":
        """Create a ZeroDBVectorStore from a list of texts."""
        store = cls(api_key=api_key, project_id=project_id, embedding=embedding, **kwargs)
        store.add_texts(texts, metadatas=metadatas)
        return store
