import pytest
from pathlib import Path

pytest.importorskip("llama_index.core", reason="Install with `uv sync --extra rag` to run LlamaIndex tests.")

from ketju.rag.llama_index.variants.simple import (
    LlamaIndexRag,
    SimplePdfRag,
    create_simple_llama_index_pdf_rag,
)
from ketju.rag.llama_index.variants.pgvector import PgVectorRag, build_pgvector_rag
from ketju.rag.llama_index.contracts import (
    LlamaIndexIngestionHandler,
    LlamaIndexQueryHandler,
    LlamaIndexStorageBackend,
)


@pytest.fixture
def docs_dir() -> Path:
    return Path("docs/interdimensional-systems")


class _DummyStorageBackend(LlamaIndexStorageBackend):
    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 _DummyIngestionHandler(LlamaIndexIngestionHandler):
    def ingest(self, rag: LlamaIndexRag, *, files=None, max_files=None) -> int:
        return 0


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


def test_simple_pdf_rag(docs_dir: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(LlamaIndexRag, "_setup_index", lambda self, storage_context=None, nodes=None: object())

    srag = LlamaIndexRag(
        name="simple_pdf_rag_test",
        docs_dir=docs_dir,
        persist_dir=tmp_path,
        storage_backend=_DummyStorageBackend(),
        ingestion_handler=_DummyIngestionHandler(),
        query_handler=_DummyQueryHandler(),
    )

    assert srag.name == "simple_pdf_rag_test"
    assert srag.persist_embedding_dir == tmp_path / "text-embedding-3-small"

def test_simple_factory_uses_persist_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(LlamaIndexRag, "_setup_index", lambda self, storage_context=None, nodes=None: object())
    rag = create_simple_llama_index_pdf_rag(name="simple_factory", persist_dir=tmp_path)

    assert isinstance(rag, LlamaIndexRag)
    assert rag.persist_embedding_dir == tmp_path / "text-embedding-3-small"


def test_simple_pdf_variant_class(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(LlamaIndexRag, "_setup_index", lambda self, storage_context=None, nodes=None: object())

    rag = SimplePdfRag(
        name="simple_pdf_variant",
        persist_dir=tmp_path,
        storage_backend=_DummyStorageBackend(),
        ingestion_handler=_DummyIngestionHandler(),
        query_handler=_DummyQueryHandler(),
    )
    assert isinstance(rag, LlamaIndexRag)
    assert rag.name == "simple_pdf_variant"


def test_pgvector_factory_wires_storage_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    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 = build_pgvector_rag(
        name="pgvector_factory",
        database_url="postgresql://user:pass@localhost:5432/vector_db",
        persist_dir=tmp_path,
    )
    assert isinstance(rag, LlamaIndexRag)
    assert rag.name == "pgvector_factory"


def test_pgvector_variant_class_wires_storage_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    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",
        database_url="postgresql://user:pass@localhost:5432/vector_db",
        persist_dir=tmp_path,
    )
    assert isinstance(rag, LlamaIndexRag)
    assert rag.name == "pgvector_variant"
