"""
RAG Comparison Example

Compares two RAG strategies side-by-side.

Requires:
  pip install 'ketju[rag]'
"""

from pathlib import Path

from ketju.config import DOCS_DIR, DATA_DIR


def _pick_demo_pdfs(max_files: int = 1) -> list[Path]:
    pdfs = list(DOCS_DIR.glob("**/*.pdf"))
    return pdfs[:max_files]


def compare_variants() -> None:
    try:
        from ketju.rag.base import BaseRAG
        from ketju.rag.llama_index.variants.simple import SimplePdfRag
        from ketju.rag.llama_index.variants.full_doc_in_ctx import FullDocInCtxRag
    except ModuleNotFoundError as e:
        raise SystemExit(
            "RAG dependencies are not installed.\n"
            "Install with: pip install 'ketju[rag]'\n"
            "Or for everything needed to run examples: pip install 'ketju[examples]'"
        ) from e

    files = _pick_demo_pdfs(max_files=1)
    if not files:
        raise SystemExit(f"No PDFs found under {DOCS_DIR}. Add a PDF under docs/ and rerun.")

    variants: dict[str, BaseRAG] = {
        "simple_pdf": SimplePdfRag(name="simple_pdf", persist_dir=DATA_DIR / "simple_pdf"),
        "full_doc_in_ctx": FullDocInCtxRag(name="full_doc_in_ctx", persist_dir=DATA_DIR / "full_doc_in_ctx"),
    }

    print("=" * 60)
    print("Ingesting")
    print("=" * 60)
    for name, rag in variants.items():
        print(f"- {name}: ingesting {len(files)} file(s)")
        rag.ingest(files=files, max_files=1)

    question = "What is this document about?"
    print("\n" + "=" * 60)
    print(f"Question: {question}")
    print("=" * 60)

    for name, rag in variants.items():
        print(f"\n--- {name} ---")
        try:
            print(rag.query(question))
        except Exception as e:
            print(f"Error: {e}")


def main() -> None:
    compare_variants()


if __name__ == "__main__":
    main()

