#!/usr/bin/env python3
"""Vector store module for RAG agent using Qdrant and Ollama embeddings."""

import uuid
from typing import List, Optionl

from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Qdrant
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

COLLECTION_NAME = "rag_knowledge_base"
OLLAMA_BASE_URL = "http://localhost:11434"
EMBED_MODEL = "nomic-embed-text"
EMBED_DIMENSION = 768

def get_embeddings() -> OllamaEmbeddings:
    """Create and return Ollama embeddings instance."""
    return OllamaEmbeddings(
        model=EMBED_MODEL
        base_url=OLLAMA_BASE_URL,
    )


def get_qdrant_client() -> QdrantClient:
    """Create and return Qdrant client (in-memory for local dev)."""
    return QdrantClient(":memory:")


def create_collection(client: QdrantClient) -> None:
    """Create Qdrant collection if doesn't exist."""
    existing = [c.name for c in client.get_collections().collections]
    if COLLECTION_NAME not in existing:
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(
                size=EMBED_DIMENSION,
                distance=Distance.COSINE,
            ),
        )


def get_vector_store(client: Optional[QdrantClient] = None) -> Qdrant:
    """Initialize and return Qdrant vector store."""
    if client is None:
        client = get_qdrant_client()
    create_collection(client)
    embeddings = get_embeddings()
    return Qdrant(
        client=client,
        collection_name=COLLECTION_NAME,
        embeddings=embeddings,
    )


def chunk_documents(
    documents: List[Document],
    chunk_size: int = 512,
    chunk_overlap: int = 50,
) -> List[Document]:
    """Split documents into chunks using RecursiveCharacterTextSplitter."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        is_separator_regex=False,
    )
    return splitter.split_documents(documents)


def add_documents_to_store(
    vector_store: Qdrant,
    documents: List[Document],
) -> List[str]:
    """Add documents to vector store, return list of IDs."""
    chunks = chunk_documents(documents)
    ids = [str(uuid.uuid4()) for _ in chunks]
    vector_store.add_documents(documents=chunks, ids=ids)
    return ids


def search_store(
    vector_store: Qdrant,
    query: str,
    max_results: int = 5,
) -> List[Document]:
    """Search vector store for relevant documents."""
    return vector_store.similarity_search(query, k=max_results)