# pyright: reportMissingTypeStubs=false

from ketju.rag.base import BaseRAG
from typing import Callable, Any
from pathlib import Path
from pypdf import PdfReader
from ketju.core.utils.files import resolve_pdf_files

try:
    from llama_index.core import Settings
except ModuleNotFoundError:
    Settings = None  # type: ignore[assignment]


class FullDocInCtxRag(BaseRAG):
    def __init__(
        self,
        *args: Any,
        llm_complete: Callable[[str], Any] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.files: list[Path] = []

        if llm_complete is not None:
            self._llm_complete = llm_complete
        elif Settings is not None:
            self._llm_complete = Settings.llm.complete
        else:  # pragma: no cover
            raise ModuleNotFoundError(
                "No LLM configured for FullDocInCtxRag. "
                "Install `ketju[rag]` for LlamaIndex defaults or pass llm_complete=..."
            )

    def _ingest(
        self,
        files: list[Path | str] | None = None,
        max_files: int | None = None,
    ) -> int:
        resolved_files = resolve_pdf_files(
            folder=self.docs_dir,
            files=files,
            max_files=max_files,
        )
        self.files.extend(resolved_files)
        return len(self.files)

    def _query(self, question: str, **kwargs: Any) -> Any:
        ctx = ""
        for file in self.files:
            reader = PdfReader(file)
            for page in reader.pages:
                ctx += (page.extract_text() or "") + "\n\n"
        return self._llm_complete(f"Question: {question}\nContext: {ctx}")
