"""
RAG (Retrieval-Augmented Generation) via ChromaDB.
Each knowledge base maps to a ChromaDB collection.
Supports document-level deletion by filename metadata.
"""

import os
from uuid import uuid4

import chromadb

from backend.config import CHROMADB_PATH

os.makedirs(CHROMADB_PATH, exist_ok=True)

_chroma_client = chromadb.PersistentClient(path=CHROMADB_PATH)


def _col_name(collection_id: str) -> str:
    """Sanitise for ChromaDB collection names (3-63 chars, alphanum/dash/underscore)."""
    name = f"kb-{collection_id}"[:63]
    return name


def create_collection(collection_id: str):
    _chroma_client.get_or_create_collection(name=_col_name(collection_id))


def delete_collection(collection_id: str):
    try:
        _chroma_client.delete_collection(name=_col_name(collection_id))
    except Exception:
        pass


def add_documents(
    collection_id: str,
    documents: list[str],
    metadatas: list[dict] | None = None,
):
    col = _chroma_client.get_or_create_collection(name=_col_name(collection_id))
    ids = [str(uuid4()) for _ in documents]
    batch_size = 500
    for i in range(0, len(documents), batch_size):
        batch_docs = documents[i : i + batch_size]
        batch_ids = ids[i : i + batch_size]
        batch_meta = metadatas[i : i + batch_size] if metadatas else None
        col.add(documents=batch_docs, ids=batch_ids, metadatas=batch_meta)


def delete_document_chunks(collection_id: str, filename: str) -> int:
    """
    Delete all vector chunks belonging to a specific document (by filename).
    Returns the number of chunks removed.
    """
    try:
        col = _chroma_client.get_collection(name=_col_name(collection_id))
    except Exception:
        return 0

    if col.count() == 0:
        return 0

    # Fetch all chunk IDs that match this filename
    try:
        results = col.get(
            where={"filename": filename},
            include=[],  # We only need IDs
        )
        chunk_ids = results.get("ids", [])
        if not chunk_ids:
            return 0

        # Delete in batches (ChromaDB can handle large deletes but let's be safe)
        batch_size = 500
        for i in range(0, len(chunk_ids), batch_size):
            batch = chunk_ids[i : i + batch_size]
            col.delete(ids=batch)

        return len(chunk_ids)
    except Exception:
        return 0


def query(collection_id: str, query_text: str, n_results: int = 8) -> str | None:
    """
    Return a formatted string of the top matching chunks,
    or None if the collection is empty / missing.
    """
    try:
        col = _chroma_client.get_collection(name=_col_name(collection_id))
    except Exception:
        return None

    if col.count() == 0:
        return None

    results = col.query(query_texts=[query_text], n_results=min(n_results, col.count()))

    docs = results.get("documents", [[]])[0]
    metas = results.get("metadatas", [[]])[0]

    if not docs:
        return None

    parts: list[str] = []
    for i, (doc, meta) in enumerate(zip(docs, metas), 1):
        src = meta.get("filename", "unknown") if meta else "unknown"
        parts.append(f"[Source {i}: {src}]\n{doc}")

    return "\n\n---\n\n".join(parts)