from __future__ import annotations

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Optional
from ketju.core.instrumentation import start_span
from ketju.config import DATA_DIR as BASE_DATA_DIR, DOCS_DIR

class BaseRAG(ABC):
    """Base class for RAG implementations.

    This base class is intentionally lightweight and has **no backend dependencies**.
    Concrete variants (e.g. LlamaIndex-based) should live in separate modules/extras.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        docs_dir: Path | None = None,
        persist_dir: Path | None = None,
        embedding_model: str = "text-embedding-3-small",
    ) -> None:
        # Backwards-compatible alias: some variants historically used `variant_name=...`.
        self.name = name or self.__class__.__name__
        self.embedding_model = embedding_model
        # Storage layout: data/<variant>/{cache,index,...}
        self.persist_dir = (persist_dir) or (BASE_DATA_DIR / (self.name).replace(" ", "_").strip())
        self.docs_dir = docs_dir or DOCS_DIR
        self.persist_dir.mkdir(parents=True, exist_ok=True)

    @property
    def persist_embedding_dir(self) -> Path:
        return self.persist_dir / self.embedding_model.replace(" ", "_").strip()

    @property
    def data_dir(self) -> Path:
        return self.persist_embedding_dir

    @property
    def cache_dir(self) -> Path:
        path = self.persist_embedding_dir / "cache"
        path.mkdir(parents=True, exist_ok=True)
        return path

    @property
    def index_dir(self) -> Path:
        path = self.persist_embedding_dir / "index"
        path.mkdir(parents=True, exist_ok=True)
        return path

    def ingest(self, *args: Any, **kwargs: Any) -> int:
        with start_span(
            f"ketju.rag.{self.name}.ingest",
        ):
            return self._ingest(*args, **kwargs)
    
    @abstractmethod
    def _ingest(self, *args: Any, **kwargs: Any) -> int:
        """Ingest documents into the RAG backend."""

    def query(self, question: str, **kwargs: Any) -> Any:
        with start_span(
            f"ketju.rag.{self.name}.query",
        ):
            return self._query(question, **kwargs)

    @abstractmethod
    def _query(self, question: str, **kwargs: Any) -> Any:
        """Query the RAG system."""

    def get_stats(self) -> dict[str, Any]:
        """Lightweight statistics about this instance."""

        return {
            "name": self.name,
            "data_dir": str(self.data_dir),
            "docs_dir": str(self.docs_dir),
        }