"""Vector store management for semantic context search.

Supports two FAISS-backed use cases:

- plan-scoped semantic search used by :class:`ContextService`
- ACMS vector indexing/query adapters used by the backend protocols
"""

from __future__ import annotations

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

from langchain_community.embeddings import FakeEmbeddings

try:
    from langchain_community.vectorstores.faiss import FAISS
except ImportError:  # pragma: no cover - exercised via graceful fallback tests
    FAISS = None

if TYPE_CHECKING:
    from langchain_core.embeddings import Embeddings

from cleveragents.application.services.config_service import ConfigService
from cleveragents.config.settings import Settings
from cleveragents.core.exceptions import ConfigurationError
from cleveragents.domain.models.core import Context
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork


class _FaissPlanStoreProtocol(Protocol):
    def similarity_search_with_score(
        self,
        query: str,
        *,
        k: int,
    ) -> list[tuple[Any, Any]]:
        _ = (query, k)
        raise NotImplementedError

    def save_local(self, folder_path: str) -> None:
        _ = folder_path
        raise NotImplementedError


class _FaissDocstoreProtocol(Protocol):
    def search(self, key: str) -> Any:
        _ = key
        raise NotImplementedError


class _FaissAcmsStoreProtocol(_FaissPlanStoreProtocol, Protocol):
    index_to_docstore_id: dict[int, str]
    docstore: _FaissDocstoreProtocol

    def add_embeddings(
        self,
        text_embeddings: list[tuple[str, list[float]]],
        *,
        metadatas: list[dict[str, str]],
        ids: list[str],
    ) -> Any:
        _ = (text_embeddings, metadatas, ids)
        raise NotImplementedError

    def similarity_search_with_score_by_vector(
        self,
        embedding: list[float],
        *,
        k: int,
        filter: Any = None,
        fetch_k: int,
    ) -> list[tuple[Any, Any]]:
        _ = (embedding, k, filter, fetch_k)
        raise NotImplementedError

    def delete(self, *, ids: list[str]) -> Any:
        _ = ids
        raise NotImplementedError

    def save_local(self, folder_path: str) -> None:
        _ = folder_path
        raise NotImplementedError


class VectorStoreService:
    """Manage creation and querying of LangChain vector stores."""

    def __init__(
        self,
        settings: Settings,
        unit_of_work: UnitOfWork,
        config_service: ConfigService | None = None,
    ) -> None:
        self.settings = settings
        self.unit_of_work = unit_of_work
        self._config_service = config_service or ConfigService()
        self._cache: dict[int, _FaissPlanStoreProtocol] = {}
        self._acms_store: _FaissAcmsStoreProtocol | None = None

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def is_enabled(self) -> bool:
        """Whether semantic vector search is enabled via settings."""

        return bool(self.settings.vector_store_enabled)

    def invalidate(self, plan_id: int | None) -> None:
        """Drop any cached vector store for the supplied plan."""

        if plan_id is None:
            return
        self._cache.pop(plan_id, None)

    def refresh_for_plan(self, plan_id: int) -> int:
        """Rebuild and persist the vector store for a plan.

        Returns the number of documents that were indexed.
        """

        self._ensure_enabled()
        documents, metadata = self._prepare_documents(plan_id)
        if not documents:
            self.invalidate(plan_id)
            self._remove_local_index(plan_id)
            return 0

        embeddings = self._create_embeddings()
        faiss_cls = self._require_faiss()
        vector_store = faiss_cls.from_texts(
            documents,
            embedding=embeddings,
            metadatas=metadata,
        )
        self._cache[plan_id] = vector_store
        vector_store.save_local(
            str(self._plan_store_dir(plan_id)),
        )
        return len(documents)

    def search(
        self,
        plan_id: int,
        query: str,
        *,
        top_k: int = 5,
        refresh_if_missing: bool = True,
    ) -> list[dict[str, Any]]:
        """Run a similarity search against the plan's vector index."""

        self._ensure_enabled()
        query = query.strip()
        if not query:
            return []

        store = self._cache.get(plan_id) or self._load_local_index(plan_id)
        if store is None and refresh_if_missing:
            if self.refresh_for_plan(plan_id) == 0:
                return []
            store = self._cache.get(plan_id) or self._load_local_index(plan_id)

        if store is None:
            return []

        limit = max(1, top_k)
        results = store.similarity_search_with_score(query, k=limit)
        formatted: list[dict[str, Any]] = []
        for document, score in results:
            formatted.append(
                {
                    "path": document.metadata.get("path"),
                    "score": float(score),
                    "snippet": document.page_content[:500],
                }
            )
        return formatted

    # ------------------------------------------------------------------
    # ACMS vector index API
    # ------------------------------------------------------------------
    def acms_backend_name(self) -> str:
        """Return the configured ACMS vector backend name."""

        raw = self._config_service.resolve("index.vector.backend").value
        return str(raw or "none").strip().lower()

    def acms_enabled(self) -> bool:
        """Whether ACMS vector indexing is configured to use FAISS."""

        return self.acms_backend_name() == "faiss"

    def acms_count(self, *, project: str | None = None) -> int:
        """Return the number of indexed ACMS embeddings."""

        store = self._acms_store or self._load_acms_store()
        if store is None:
            return 0

        if project is None:
            return len(getattr(store, "index_to_docstore_id", {}))

        count = 0
        for _store_id, metadata in self._iter_store_metadata(store):
            if metadata.get("project") == project:
                count += 1
        return count

    def acms_index_embedding(
        self,
        project: str,
        doc_id: str,
        embedding: list[float],
        metadata: dict[str, str],
        *,
        content: str | None = None,
    ) -> None:
        """Index one ACMS embedding into the shared FAISS store."""

        if not project or not project.strip():
            raise ValueError("project must be a non-empty string")
        if not doc_id or not doc_id.strip():
            raise ValueError("doc_id must be a non-empty string")
        if not embedding:
            raise ValueError("embedding must be a non-empty list")
        if not self.acms_enabled():
            raise ConfigurationError(
                "ACMS vector indexing is disabled. Set index.vector.backend to faiss "
                "to enable the FAISS ACMS backend."
            )

        faiss_cls = self._require_faiss()
        store = self._acms_store or self._load_acms_store()
        stored_metadata = dict(metadata)
        stored_metadata["project"] = project
        stored_metadata["doc_id"] = doc_id
        payload = (content or "").strip() or stored_metadata.get(
            "location", f"embedding:{doc_id}"
        )
        store_id = self._acms_store_id(project, doc_id)

        if store is None:
            store = cast(
                _FaissAcmsStoreProtocol,
                faiss_cls.from_embeddings(
                    [(payload, list(embedding))],
                    self._create_acms_embeddings(),
                    metadatas=[stored_metadata],
                    ids=[store_id],
                ),
            )
        else:
            self._delete_store_id(store, store_id)
            store.add_embeddings(
                [(payload, list(embedding))],
                metadatas=[stored_metadata],
                ids=[store_id],
            )

        self._acms_store = store
        self._save_acms_store()

    def acms_search_by_vector(
        self,
        embedding: list[float],
        *,
        top_k: int = 20,
        scope: frozenset[str] | None = None,
        project: str | None = None,
    ) -> list[dict[str, Any]]:
        """Search ACMS embeddings by vector similarity."""

        if not embedding:
            raise ValueError("embedding must be a non-empty list")
        if top_k < 1:
            raise ValueError(f"top_k must be positive, got {top_k}")
        if not self.acms_enabled():
            return []

        store = self._acms_store or self._load_acms_store()
        if store is None:
            return []

        scope_set = set(scope or ())

        def _matches(metadata: dict[str, Any]) -> bool:
            if project is not None and metadata.get("project") != project:
                return False
            return not scope_set or metadata.get("resource_id") in scope_set

        raw_results = store.similarity_search_with_score_by_vector(
            list(embedding),
            k=top_k,
            filter=_matches if (scope_set or project is not None) else None,
            fetch_k=max(20, top_k * 5),
        )
        results: list[dict[str, Any]] = []
        for document, distance in raw_results:
            metadata = {
                str(key): str(value)
                for key, value in document.metadata.items()
                if value is not None
            }
            results.append(
                {
                    "content": document.page_content,
                    "distance": float(distance),
                    "score": self._distance_to_relevance(distance),
                    "metadata": metadata,
                }
            )
        return results[:top_k]

    def acms_remove_embedding(self, project: str, doc_id: str) -> None:
        """Remove one ACMS embedding from the shared FAISS store."""

        if not project or not project.strip():
            raise ValueError("project must be a non-empty string")
        if not doc_id or not doc_id.strip():
            raise ValueError("doc_id must be a non-empty string")

        store = self._acms_store or self._load_acms_store()
        if store is None:
            return

        self._delete_store_id(store, self._acms_store_id(project, doc_id))
        if getattr(store, "index_to_docstore_id", {}):
            self._acms_store = store
            self._save_acms_store()
            return

        self._acms_store = None
        self._remove_acms_store_files()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _ensure_enabled(self) -> None:
        if not self.is_enabled():
            raise ConfigurationError(
                "Vector store support is disabled. Set "
                "CLEVERAGENTS_VECTOR_STORE_ENABLED to true to enable semantic search."
            )

    def _require_faiss(self) -> Any:
        if FAISS is None:
            raise ConfigurationError(
                "FAISS library not installed -- vector search will not work"
            )
        return FAISS

    def _prepare_documents(
        self, plan_id: int
    ) -> tuple[list[str], list[dict[str, Any]]]:
        with self.unit_of_work.transaction() as ctx:
            contexts: Iterable[Context] = ctx.contexts.get_for_plan(plan_id)

        documents: list[str] = []
        metadata: list[dict[str, Any]] = []
        for context in contexts:
            if not context.content:
                continue
            documents.append(context.content)
            metadata.append({"path": context.path, "plan_id": plan_id})
        return documents, metadata

    def _create_embeddings(self) -> Embeddings:
        provider = (self.settings.vector_embeddings_provider or "fake").lower()
        if provider == "fake":
            return FakeEmbeddings(size=self.settings.vector_embeddings_dimension)
        if provider in {"openai", "azure", "azureopenai"}:
            try:
                from langchain_openai import OpenAIEmbeddings
            except ImportError as exc:  # pragma: no cover - import guard
                raise ConfigurationError(
                    "langchain-openai is required for OpenAI embeddings"
                ) from exc

            model = self.settings.vector_embeddings_model or "text-embedding-3-small"
            return OpenAIEmbeddings(model=model)

        raise ConfigurationError(
            "Unsupported vector embeddings provider "
            f"'{self.settings.vector_embeddings_provider}'"
        )

    def _create_acms_embeddings(self) -> Embeddings:
        provider = (
            str(
                self._config_service.resolve("index.embedding.provider").value or "fake"
            )
            .strip()
            .lower()
        )
        model_value = self._config_service.resolve("index.embedding.model").value
        dims_value = self._config_service.resolve("index.embedding.dimensions").value

        if provider in {"fake", "local"}:
            if dims_value is None:
                dims_value = self.settings.vector_embeddings_dimension
            return FakeEmbeddings(size=int(dims_value))

        if provider in {"openai", "azure", "azureopenai"}:
            try:
                from langchain_openai import OpenAIEmbeddings
            except ImportError as exc:  # pragma: no cover - import guard
                raise ConfigurationError(
                    "langchain-openai is required for OpenAI embeddings"
                ) from exc

            kwargs: dict[str, Any] = {"model": model_value or "text-embedding-3-small"}
            if dims_value is not None:
                kwargs["dimensions"] = int(dims_value)
            return OpenAIEmbeddings(**kwargs)

        raise ConfigurationError(f"Unsupported ACMS embedding provider '{provider}'")

    def _plan_store_dir(self, plan_id: int) -> Path:
        base = self.settings.resolve_vector_store_path()
        plan_dir = base / f"plan_{plan_id}"
        plan_dir.mkdir(parents=True, exist_ok=True)
        return plan_dir

    def _load_local_index(self, plan_id: int) -> _FaissPlanStoreProtocol | None:
        plan_dir = self._plan_store_dir(plan_id)
        index_path = plan_dir / "index.faiss"
        store_path = plan_dir / "index.pkl"
        if not index_path.exists() or not store_path.exists():
            return None
        embeddings = self._create_embeddings()
        faiss_cls = self._require_faiss()
        try:
            store = faiss_cls.load_local(
                str(plan_dir),
                embeddings,
                allow_dangerous_deserialization=True,
            )
        except (ValueError, FileNotFoundError):
            return None
        self._cache[plan_id] = store
        return store

    def _remove_local_index(self, plan_id: int) -> None:
        plan_dir = self._plan_store_dir(plan_id)
        index_path = plan_dir / "index.faiss"
        store_path = plan_dir / "index.pkl"
        for path in (index_path, store_path):
            if path.exists():
                path.unlink(missing_ok=True)

    def _acms_store_dir(self) -> Path:
        raw_dir = self._config_service.resolve("index.vector.dir").value
        base = Path(str(raw_dir)).expanduser()
        if not base.is_absolute():
            base = Path.cwd() / base
        base.mkdir(parents=True, exist_ok=True)
        return base.resolve()

    def _load_acms_store(self) -> _FaissAcmsStoreProtocol | None:
        store_dir = self._acms_store_dir()
        index_path = store_dir / "index.faiss"
        store_path = store_dir / "index.pkl"
        if not index_path.exists() or not store_path.exists():
            return None

        faiss_cls = self._require_faiss()
        try:
            store = faiss_cls.load_local(
                str(store_dir),
                self._create_acms_embeddings(),
                allow_dangerous_deserialization=True,
            )
        except (FileNotFoundError, ValueError):
            return None
        self._acms_store = store
        return store

    def _save_acms_store(self) -> None:
        if self._acms_store is None:
            return
        self._acms_store.save_local(str(self._acms_store_dir()))

    def _remove_acms_store_files(self) -> None:
        store_dir = self._acms_store_dir()
        for path in (store_dir / "index.faiss", store_dir / "index.pkl"):
            if path.exists():
                path.unlink(missing_ok=True)

    def _acms_store_id(self, project: str, doc_id: str) -> str:
        return f"{project}::{doc_id}"

    def _delete_store_id(self, store: _FaissAcmsStoreProtocol, store_id: str) -> None:
        docstore_ids = set(getattr(store, "index_to_docstore_id", {}).values())
        if store_id in docstore_ids:
            store.delete(ids=[store_id])

    def _iter_store_metadata(
        self, store: _FaissAcmsStoreProtocol
    ) -> Iterable[tuple[str, dict[str, str]]]:
        for store_id in getattr(store, "index_to_docstore_id", {}).values():
            document = store.docstore.search(store_id)
            metadata = getattr(document, "metadata", {})
            yield store_id, {str(key): str(value) for key, value in metadata.items()}

    def _distance_to_relevance(self, distance: Any) -> float:
        numeric_distance = max(float(distance), 0.0)
        return 1.0 / (1.0 + numeric_distance)
