from __future__ import annotations

from pathlib import Path

import pytest


def test_full_doc_in_ctx_allows_custom_llm(tmp_path: Path) -> None:
    pytest.importorskip("pypdf", reason="Install rag extras to run this test")
    from ketju.rag.llama_index.variants.full_doc_in_ctx import FullDocInCtxRag

    calls: list[str] = []

    def _fake_complete(prompt: str) -> str:
        calls.append(prompt)
        return "ok"

    rag = FullDocInCtxRag(
        name="full_doc_test",
        docs_dir=tmp_path,
        persist_dir=tmp_path,
        llm_complete=_fake_complete,
    )
    ingested = rag.ingest(files=[], max_files=1)
    answer = rag.query("What is this?")

    assert ingested == 0
    assert answer == "ok"
    assert calls and "Question: What is this?" in calls[0]


def test_docling_variants_use_llamaindex_rag_base(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    pytest.importorskip("llama_index.core", reason="Needs llama_index.core")
    pytest.importorskip("llama_index.readers.docling", reason="Install docling extras to run this test")
    pytest.importorskip("llama_index.node_parser.docling", reason="Install docling extras to run this test")

    from ketju.rag.llama_index.variants.simple import LlamaIndexRag
    from ketju.rag.llama_index.variants.simple_docling import SimpleDoclingRag
    from ketju.rag.llama_index.variants.simple_docling_full_ctx import SimpleDoclingFullCtxRag

    class _DummyStorageBackend:
        def get_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> object:
            return object()

        def create_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> object:
            return object()

        def save(self, *, vector_index: object, index_dir: Path, data_dir: Path) -> None:
            return None

    class _DummyQueryHandler:
        def query(self, rag: LlamaIndexRag, question: str, **kwargs: object) -> str:
            return "ok"

    monkeypatch.setattr(LlamaIndexRag, "_setup_index", lambda self, storage_context=None, nodes=None: object())

    docling_rag = SimpleDoclingRag(
        name="docling",
        persist_dir=tmp_path,
        storage_backend=_DummyStorageBackend(),
        query_handler=_DummyQueryHandler(),
    )
    full_ctx_rag = SimpleDoclingFullCtxRag(
        name="docling_full_ctx",
        persist_dir=tmp_path,
        storage_backend=_DummyStorageBackend(),
        query_handler=_DummyQueryHandler(),
    )

    assert isinstance(docling_rag, LlamaIndexRag)
    assert isinstance(full_ctx_rag, LlamaIndexRag)


def test_chapter_rag_init_without_external_calls(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    pytest.importorskip("chromadb", reason="Install rag extras to run this test")
    pytest.importorskip("llama_index.vector_stores.chroma", reason="Install rag extras to run this test")
    pytest.importorskip("langchain_docling", reason="Install docling extras to run this test")

    from ketju.rag.llama_index.variants.chapter import ChapterRag

    monkeypatch.setattr(ChapterRag, "_setup_vector_store", lambda self: None)
    monkeypatch.setattr(ChapterRag, "_load_chapters", lambda self: setattr(self, "chapters", {}))

    rag = ChapterRag(name="chapter_rag_test", docs_dir=tmp_path, persist_dir=tmp_path)
    assert rag.name == "chapter_rag_test"


def test_pgvector_variant_class_wiring(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    pytest.importorskip("llama_index.core", reason="Needs llama_index.core")

    from ketju.rag.llama_index.variants.simple import LlamaIndexRag
    from ketju.rag.llama_index.variants.pgvector import PgVectorRag
    import ketju.rag.llama_index.variants.pgvector as pgvector_mod

    class _FakePgBackend:
        def __init__(self, **kwargs: object) -> None:
            self.kwargs = kwargs

        def get_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> object:
            return object()

        def create_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> object:
            return object()

        def save(self, *, vector_index: object, index_dir: Path, data_dir: Path) -> None:
            return None

    monkeypatch.setattr(pgvector_mod, "PgVectorStorageBackend", _FakePgBackend)
    monkeypatch.setattr(LlamaIndexRag, "_setup_index", lambda self, storage_context=None, nodes=None: object())

    rag = PgVectorRag(
        name="pgvector_variant_in_wiring_test",
        database_url="postgresql://user:pass@localhost:5432/vector_db",
        persist_dir=tmp_path,
    )
    assert isinstance(rag, LlamaIndexRag)
    assert rag.name == "pgvector_variant_in_wiring_test"
