from __future__ import annotations

from pathlib import Path
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Protocol

if TYPE_CHECKING:  # pragma: no cover
    from llama_index.core import StorageContext, VectorStoreIndex
    from .variants.simple import LlamaIndexRag
    from llama_index.core.base.response.schema import RESPONSE_TYPE


class LlamaIndexStorageBackend(Protocol):
    """Internal storage strategy contract for LlamaIndex RAG implementations."""

    def get_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> StorageContext: ...

    def create_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> StorageContext: ...

    def save(self, *, vector_index: VectorStoreIndex, index_dir: Path, data_dir: Path) -> None: ...


class LlamaIndexIngestionHandler(Protocol):
    """Internal ingestion strategy contract for LlamaIndex RAG implementations."""

    def ingest(
        self,
        rag: LlamaIndexRag,
        *,
        files: Sequence[Path | str] | None = None,
        max_files: int | None = None,
    ) -> int: ...

class LlamaIndexQueryHandler(Protocol):
    """Internal query strategy contract for LlamaIndex RAG implementations."""

    def query(self, rag: LlamaIndexRag, question: str, **kwargs: Any) -> RESPONSE_TYPE: ...