from pathlib import Path
from uuid import uuid4

from langchain_chroma import Chroma, Settings
from langchain_ollama import OllamaEmbeddings
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

CHROMA_DIR = "./chroma_db"
COLLECTION_NAME = "local_kb"
EMBED_MODEL = "nomic-embed-text"
OLLAMA_BASE_URL = "http://127.0.0.1:11434"


def create_vectorstore(persist_directory: str = CHROMA_DIR):
    """Return a Chroma collection configured for the local knowledge base.

    The collection is persisted in ``persist_directory`` and named
    ``COLLECTION_NAME``.
    """
    client = Chroma(Settings(persist_directory=persist_directory))
    collection = client.get_or_create_collection(name=COLLECTION_NAME)
    return collection


def load_documents(directory: str, collection) -> int:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        add_start_index=True,
    )

    base_path = Path(directory)
    if not base_path.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")

    raw_docs: list[Document] = []
    for pattern in ("*.txt", "*.md"):
        for path in sorted(base_path.rglob(pattern)):
            content = path.read_text(encoding="utf-8")
            raw_docs.append(
                Document(
                    page_content=content,
                    metadata={"source": str(path)},
                )
            )

    if not raw_docs:
        return 0

    chunks = splitter.split_documents(raw_docs)
    ids = [str(uuid4()) for _ in chunks]
    embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_BASE_URL)
    vectors = embeddings.embed_documents([doc.page_content for doc in chunks])
    collection.add(
        ids=ids,
        documents=[doc.page_content for doc in chunks],
        embeddings=vectors,
        metadatas=[doc.metadata for doc in chunks],
    )
    return len(chunks)