from __future__ import annotations

from pathlib import Path
from collections.abc import Sequence
from typing import Optional, TYPE_CHECKING
from ketju.core.utils.files import resolve_pdf_files

from ketju.rag.base import BaseRAG
from ketju.config import DOCS_DIR
from ketju.rag.llama_index.contracts import LlamaIndexStorageBackend
from ketju.core.instrumentation import start_span
from ketju.rag.llama_index.contracts import LlamaIndexIngestionHandler, LlamaIndexQueryHandler
from typing import Any

if TYPE_CHECKING:
    from ketju.rag.llama_index.storage.local import LocalStorageBackend
    from ketju.rag.llama_index.ingestion.simple_pdf import SimplePdfIngestionHandler
    from ketju.rag.llama_index.query.simple import SimpleLlamaIndexQueryHandler

try:
    from llama_index.core import Settings, StorageContext, VectorStoreIndex
    from llama_index.core.schema import Node
    from llama_index.core.storage.docstore.types import BaseDocumentStore
    from llama_index.core.vector_stores.types import BasePydanticVectorStore
    from ketju.rag.llama_index.storage.local import LocalStorageBackend
    from ketju.rag.llama_index.ingestion.simple_pdf import SimplePdfIngestionHandler
    from ketju.rag.llama_index.query.simple import SimpleLlamaIndexQueryHandler
except ModuleNotFoundError as e:  # pragma: no cover
    raise ModuleNotFoundError(
        "The LlamaIndex backend is not installed. Install it with:\n"
        "  pip install 'ketju[rag]'\n"
        "Or, for examples:\n"
        "  pip install 'ketju[examples]'"
    ) from e


class LlamaIndexRag(BaseRAG):
    def __init__(self,
        name: Optional[str] = None,
        *,
        docs_dir: Path = DOCS_DIR or None,
        persist_dir: Path | None = None,
        embedding_model: str = "text-embedding-3-small",
        storage_backend: LlamaIndexStorageBackend | None = None,
        ingestion_handler: LlamaIndexIngestionHandler | None = None,
        query_handler: LlamaIndexQueryHandler | None = None,
        ) -> None:
        super().__init__(
            name=name,
            docs_dir=docs_dir,
            persist_dir=persist_dir,
            embedding_model=embedding_model,
        )

        self.storage_backend = storage_backend or LocalStorageBackend()
        self.ingestion_handler = ingestion_handler or SimplePdfIngestionHandler()
        self.query_handler = query_handler or SimpleLlamaIndexQueryHandler()
        self.vector_store_index = self._setup_index(
            storage_context=self.storage_backend.get_storage_context(
                index_dir=self.index_dir,
                data_dir=self.data_dir,
                name="default",
            )
        )

    def _setup_index(self, storage_context: Optional[StorageContext] = None, nodes: list[Node] | None = None):
        _nodes = nodes or []
        if _nodes:
            return VectorStoreIndex(_nodes, storage_context=storage_context, embed_model=Settings.embed_model)
        return VectorStoreIndex(
            [],
            storage_context=storage_context,
            embed_model=Settings.embed_model,
        )

    @property
    def vector_index(self) -> VectorStoreIndex:
        return self.vector_store_index

    @property
    def docstore(self) -> BaseDocumentStore:
        return self.vector_index.storage_context.docstore

    @property
    def vector_store(self) -> BasePydanticVectorStore:
        return self.vector_index.vector_store

    def get_storage_context(self, name: str = "default") -> StorageContext:
        return self.storage_backend.get_storage_context(
            index_dir=self.index_dir,
            data_dir=self.data_dir,
            name=name,
        )

    def _ingest(
        self,
        files: Sequence[Path | str] | None = None,
        max_files: int | None = None,
    ) -> int:
        with start_span(
            f"ketju.rag.{self.name}.ingest",
            attributes={
                "ketju.rag.variant": self.name,
                "ketju.rag.max_files": max_files if max_files is not None else -1,
            },
        ):
            return self.ingestion_handler.ingest(self, files=files, max_files=max_files)

    def _save(self) -> None:
        self.storage_backend.save(
            vector_index=self.vector_index,
            index_dir=self.index_dir,
            data_dir=self.data_dir,
        )

    def save(self) -> None:
        self._save()

    def _query(self, question: str, **kwargs: Any) -> Any:
        return self.query_handler.query(self, question, **kwargs)

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


class SimplePdfRag(LlamaIndexRag):
    """Compatibility alias for the default simple PDF configuration."""


def create_simple_llama_index_pdf_rag(
    name: str = "simple_pdf_rag",
    *,
    persist_dir: Path | None = None,
    docs_dir: Path | None = None,
    embedding_model: str = "text-embedding-3-small",
) -> LlamaIndexRag:
    return LlamaIndexRag(
        name=name,
        persist_dir=persist_dir,
        docs_dir=docs_dir or DOCS_DIR,
        embedding_model=embedding_model,
    )
