import os
from pathlib import Path
from typing import Optional

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient


def create_vectorstore(
    persist_directory: str = "./qdrant_db",
    collection_name: str = "documents",
    embedding_model: str = "nomic-embed-text",
) -> QdrantVectorStore:
    """
    Create a Qdrant vector store backed by Ollama embeddings.

    Parameters
    ----------
    persist_directory : str
        Path to the directory where Qdrant will store its data.
    collection_name : str
        Name of the collection to use.
    embedding_model : str
        Ollama model name for embeddings.

    Returns
    -------
    QdrantVectorStore
        Initialized vector store.
    """
    # Ensure the persistence directory exists
    Path(persist_directory).mkdir(parents=True, exist_ok=True)

    # Create a local Qdrant client that stores data on disk
    client = QdrantClient(path=persist_directory)

    # Initialize embeddings
    embeddings = OllamaEmbeddings(model=embedding_model)

    # Create or load the collection
    vectorstore = QdrantVectorStore(
        client=client,
        embeddings=embeddings,
        collection_name=collection_name,
    )
    return vectorstore


def load_documents(
    directory: str,
    vectorstore: QdrantVectorStore,
    chunk_size: int = 1000,
    chunk_overlap: int = 200,
) -> None:
    """
    Load all .txt and .md files from a directory, split them into chunks,
    embed, and add to the vector store.

    Parameters
    ----------
    directory : str
        Path to the directory containing documents.
    vectorstore : QdrantVectorStore
        The vector store to populate.
    chunk_size : int
        Maximum size of each chunk in characters.
    chunk_overlap : int
        Number of overlapping characters between chunks.
    """
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from langchain.docstore.document import Document

    # Prepare text splitter
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )

    # Collect all documents
    docs = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.lower().endswith((".txt", ".md")):
                file_path = os.path.join(root, file)
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()
                # Split content into chunks
                chunks = splitter.split_text(content)
                # Create Document objects with metadata
                for i, chunk in enumerate(chunks):
                    doc = Document(
                        page_content=chunk,
                        metadata={
                            "source": file_path,
                            "chunk_index": i,
                        },
                    )
                    docs.append(doc)

    if docs:
        vectorstore.add_documents(docs)
    else:
        print("No documents found to load.")


def collection_exists(vectorstore: QdrantVectorStore) -> bool:
    """
    Check if the collection already exists in Qdrant.

    Parameters
    ----------
    vectorstore : QdrantVectorStore
        The vector store to check.

    Returns
    -------
    bool
        True if collection exists, False otherwise.
    """
    try:
        vectorstore.client.get_collection(vectorstore.collection_name)
        return True
    except Exception:
        return False