from langchain_qdrant import QdrantVectorStore
from langchain_ollama import OllamaEmbeddings
from langchain_core.documents import Document
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
from langchain.text_splitter import RecursiveCharacterTextSplitter
import uuid

COLLECTION_NAME = "knowledge_base"
EMBEDDING_MODEL = "nomic-embed-text"
VECTOR_SIZE = 768


def get_embeddings():
    return OllamaEmbeddings(model=EMBEDDING_MODEL)


def get_qdrant_client():
    return QdrantClient(host="localhost", port=6333)


def init_collection(client: QdrantClient):
    collections = [c.name for c in client.get_collections().collections]
    if COLLECTION_NAME not in collections:
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
        )


def get_vector_store() -> QdrantVectorStore:
    client = get_qdrant_client()
    init_collection(client)
    embeddings = get_embeddings()
    return QdrantVectorStore(
        client=client,
        collection_name=COLLECTION_NAME,
        embedding=embeddings,
    )


def add_documents(content: str, title: str) -> int:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        separators=["\n\n", "\n", ".", " ", ""],
    )
    chunks = splitter.split_text(content)
    docs = [
        Document(
            page_content=chunk,
            metadata={"title": title, "chunk_index": i, "source": title},
        )
        for i, chunk in enumerate(chunks)
    ]
    store = get_vector_store()
    store.add_documents(docs)
    return len(docs)


def search_documents(query: str, max_results: int = 5) -> list[dict]:
    store = get_vector_store()
    results = store.similarity_search_with_relevance_scores(query, k=max_results)
    output = []
    for doc, score in results:
        output.append(
            {
                "content": doc.page_content,
                "metadata": doc.metadata,
                "score": round(score, 4),
            }
        )
    return output