"""
Qdrant vector store wrapper for adding and searching documents.
"""

from pathlib import Path
from typing import List, Optional

from qdrant_client import QdrantClient
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_qdrant import QdrantVectorStore
from langchain.docstore.document import Document

# Configuration
QDRANT_HOST = "localhost"
QDRANT_PORT = 6333
COLLECTION_NAME = "knowledge_base"
EMBEDDING_MODEL = "nomic-embed-text"

# Initialize embedding model
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

# Initialize Qdrant client
client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)

# Ensure collection exists
if COLLECTION_NAME not in client.get_collections().collections:
    client.recreate_collection(collection_name=COLLECTION_NAME, vectors_config=client.get_default_vector_config())

# Create vector store instance
vector_store = QdrantVectorStore(
    client=client,
    collection_name=COLLECTION_NAME,
    embeddings=embeddings,
)

# Text splitter
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)


def add_document(content: str, title: str) -> None:
    """Add a document to the vector store.

    Parameters
    ----------
    content: str
        Full text of the document.
    title: str
        Title or identifier for the document.
    """
    # Split into chunks
    chunks = splitter.split_text(content)
    docs = [Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
            for i, chunk in enumerate(chunks)]
    # Add to vector store
    vector_store.add_documents(docs)


def search(query: str, max_results: int = 5) -> List[Document]:
    """Semantic search in the knowledge base.

    Parameters
    ----------
    query: str
        Search query.
    max_results: int
        Number of top results to return.
    """
    return vector_store.similarity_search(query, k=max_results)


def load_directory(directory: str) -> None:
    """Load all text files from a directory into the vector store.

    Parameters
    ----------
    directory: str
        Path to directory containing .txt files.
    """
    for file_path in Path(directory).rglob("*.txt"):
        text = file_path.read_text(encoding="utf-8")
        title = file_path.stem
        add_document(text, title)

# Expose public API
__all__ = [
    "add_document",
    "search",
    "load_directory",
]