"""
RAG pipeline: load knowledge base -> chunk -> embed (bge-m3) -> Chroma -> retrieve.

Embeddings are computed explicitly via the LLM backend (Ollama/vLLM) and passed
to Chroma, so the retrieval path is fully under our control and uses the same
multilingual model for documents and queries.
"""
import os
import glob
import chromadb

from .config import settings
from .embeddings import embed_texts


def get_collection(reset: bool = False):
    db = chromadb.PersistentClient(path=settings.chroma_dir)
    if reset:
        try:
            db.delete_collection(settings.collection)
        except Exception:
            pass
    # cosine distance pairs well with the multilingual embedders we use
    return db.get_or_create_collection(settings.collection, metadata={"hnsw:space": "cosine"})


def chunk_text(text: str, size: int, overlap: int) -> list[str]:
    """Pack paragraphs into ~`size`-char chunks, then add `overlap` between them."""
    paras = [p.strip() for p in text.split("\n\n") if p.strip()]
    chunks, cur = [], ""
    for p in paras:
        if len(cur) + len(p) + 2 <= size:
            cur = (cur + "\n\n" + p).strip()
        else:
            if cur:
                chunks.append(cur)
            if len(p) <= size:
                cur = p
            else:                                  # hard-split an oversized paragraph
                for i in range(0, len(p), size - overlap):
                    chunks.append(p[i:i + size])
                cur = ""
    if cur:
        chunks.append(cur)

    if overlap > 0 and len(chunks) > 1:
        stitched = [chunks[0]]
        for i in range(1, len(chunks)):
            tail = chunks[i - 1][-overlap:]
            stitched.append((tail + "\n" + chunks[i]).strip())
        chunks = stitched
    return chunks


def ingest(kb_dir: str = "knowledge_base", reset: bool = False) -> int:
    """(Re)build the vector store from every *.md file in the knowledge base.

    Pass reset=True to drop the collection first (needed when the embedding
    model — and therefore the vector dimension — changes).
    """
    col = get_collection(reset=reset)
    ids, docs, metas = [], [], []
    for path in sorted(glob.glob(os.path.join(kb_dir, "*.md"))):
        lang = "kk" if path.endswith("_kk.md") else "ru"
        with open(path, encoding="utf-8") as f:
            text = f.read()
        for i, chunk in enumerate(chunk_text(text, settings.chunk_size, settings.chunk_overlap)):
            ids.append(f"{os.path.basename(path)}::{i}")
            docs.append(chunk)
            metas.append({"source": os.path.basename(path), "lang": lang})

    if not docs:
        return 0
    embeddings = embed_texts(docs)
    col.upsert(ids=ids, documents=docs, embeddings=embeddings, metadatas=metas)
    return len(docs)


def retrieve(query: str, top_k: int | None = None, lang: str | None = None) -> list[dict]:
    col = get_collection()
    q_emb = embed_texts([query])[0]
    where = {"lang": lang} if lang else None
    res = col.query(
        query_embeddings=[q_emb],
        n_results=top_k or settings.top_k,
        where=where,
    )
    docs = res["documents"][0]
    metas = res["metadatas"][0]
    dists = res["distances"][0]
    # cosine distance -> similarity score in [0, 1]
    return [
        {"text": d, "source": m["source"], "score": max(0.0, 1.0 - dist)}
        for d, m, dist in zip(docs, metas, dists)
    ]